Xenova HF Staff commited on 3 days ago

Commit

2157e54

verified ·

1 Parent(s): 36d2a82

sync 91d990483a17

Browse files

Files changed (20) hide show

README.md +61 -13
build/webgpu/attention-rank4-tiled.wgsl.jinja +26 -13
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +2 -1
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +31 -37
build/webgpu/attn-flash-online.wgsl.jinja +30 -60
build/webgpu/attn-flash-prefill-cluster.wgsl.jinja +172 -70
build/webgpu/attn-flash-q32-broadcast.wgsl.jinja +1 -4
build/webgpu/attn-materialized-apply-f32.wgsl.jinja +5 -7
build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja +3 -4
build/webgpu/attn-materialized-score-f32.wgsl.jinja +6 -8
build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +17 -14
build/webgpu/attn-materialized-softmax-f32.wgsl.jinja +27 -34
build/webgpu/attn-online-scalar.wgsl.jinja +6 -28
build/webgpu/attn-small-head-parallel.wgsl.jinja +0 -3
build/webgpu/bench.json +1 -2
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +67 -20
build/webgpu/mha-small-seq-blocked.wgsl.jinja +166 -0
build/webgpu/mha-small-seq.wgsl.jinja +3 -3
build/webgpu/test.json +292 -50

README.md CHANGED Viewed

@@ -12,25 +12,25 @@ tags:
 ## Description
-Computes multi-head self- or cross-attention over explicit query, key, and value tensors, with an optional fused QKV projection bias and additive attention bias. Supports causal masking through `unidirectional` and a configurable score scale that defaults to `1 / sqrt(head_size)`. The current kernels require query, key, and value to have the same hidden width; distinct value-head widths, KV-cache, key-padding-mask, and diagnostic-QK modes are not yet implemented.
 See the [ONNX Runtime `MultiHeadAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MultiHeadAttention) for the reference semantics.
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `query` | `queryT` | `T` | `3` | — | Query tensor of shape `(batch_size, sequence_length, hidden_size)`. | required |
-| `key` | `keyT` | `T` | `3` | — | Key tensor of shape `(batch_size, kv_sequence_length, hidden_size)`; the current kernels require its hidden width to equal the query hidden width. | required |
-| `value` | `valueT` | `T` | `3` | — | Value tensor of shape `(batch_size, kv_sequence_length, v_hidden_size)`; the current supported subset requires `v_hidden_size` to equal the query hidden width. | required |
-| `bias` | `biasT` | `T` | `1` | — | Optional fused bias from input projection with shape `(hidden_size + hidden_size + v_hidden_size)`. | optional |
-| `attention_bias` | `attentionBiasT` | `T` | `4` | — | Optional additive bias applied to the attention scores before softmax, of shape `(batch_size or 1, num_heads or 1, sequence_length, kv_sequence_length)`; the first two dimensions broadcast. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | `3` | `[query[0], query[1], value[2]]` | Attention output of shape `(batch_size, sequence_length, v_hidden_size)`. | required |
 ## Attributes
@@ -38,9 +38,9 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
-| `unidirectional` | `0` | When 1, applies an upper-left causal mask: query position `i` attends only to key positions `0..min(i, kv_sequence_length - 1)`. |
 | `num_heads` | — | Number of attention heads. |
 | `scale` | — | Optional score scale; zero or omission selects `1 / sqrt(head_size)`. |
 ## Type constraints
@@ -48,13 +48,56 @@ Attributes and default values (overridable per request):
 | --- | --- |
 | `T` | `float32`, `float16` |
 ## 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
@@ -71,14 +114,19 @@ Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgr
 - [`attn-materialized-softmax-f32.wgsl.jinja`](build/webgpu/attn-materialized-softmax-f32.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`attn-small-head-parallel.wgsl.jinja`](build/webgpu/attn-small-head-parallel.wgsl.jinja)
 - [`mha-small-seq.wgsl.jinja`](build/webgpu/mha-small-seq.wgsl.jinja)
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Description
+Computes multi-head self- or cross-attention over explicit query, key, and value tensors, with an optional fused QKV projection bias and additive attention bias. Supports causal masking through `unidirectional` and a configurable score scale that defaults to `1 / sqrt(head_size)`. Query, key, and value must have the same hidden width; distinct value-head widths, KV-cache, key-padding-mask, and diagnostic-QK modes are unsupported.
 See the [ONNX Runtime `MultiHeadAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MultiHeadAttention) for the reference semantics.
 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | `3` | — | Query tensor of shape `(batch_size, sequence_length, hidden_size)`. | required |
+| `keyT` | `key` | `T` | `3` | — | Key tensor of shape `(batch_size, kv_sequence_length, hidden_size)`; its hidden width must equal the query hidden width. | required |
+| `valueT` | `value` | `T` | `3` | — | Value tensor of shape `(batch_size, kv_sequence_length, v_hidden_size)`; `v_hidden_size` must equal the query hidden width. | required |
+| `biasT` | `bias` | `T` | `1` | — | Optional fused bias from input projection with shape `(hidden_size + hidden_size + v_hidden_size)`. | optional |
+| `attentionBiasT` | `attention_bias` | `T` | `4` | — | Optional additive bias applied to the attention scores before softmax, of shape `(batch_size or 1, num_heads or 1, sequence_length, kv_sequence_length)`; the first two dimensions broadcast. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `3` | `[queryT[0], queryT[1], valueT[2]]` | Attention output of shape `(batch_size, sequence_length, v_hidden_size)`. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `num_heads` | — | Number of attention heads. |
 | `scale` | — | Optional score scale; zero or omission selects `1 / sqrt(head_size)`. |
+| `unidirectional` | `0` | When 1, applies an upper-left causal mask: query position `i` attends only to key positions `0..min(i, kv_sequence_length - 1)`. |
 ## Type constraints
 | --- | --- |
 | `T` | `float32`, `float16` |
+## 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_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences with bias and optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
+- `qkv_no_bias_small_seq_blocked` — Blocked whole-head attention for short key sequences without bias and with optional causal masking. A workgroup stages one head's K and V once, then assigns each query in a small block to a group of key lanes whose register-resident online-softmax partials are merged in one pass.
+- `qkv_no_bias_small_seq` — Whole-head attention for short bidirectional float32 requests without bias. One workgroup stages a head's complete key and value planes, and each participating invocation owns one query row through scoring, normalization, and context accumulation.
+- `qkv_no_bias_tiled_nosg` — Portable tiled prefill route for devices without an admissible subgroup path. It computes attention online without materializing the full score matrix.
+- `qkv_no_bias_tiled_attn_bias_nosg` — Portable tiled prefill route for additive attention bias when no subgroup path is admissible. It computes attention online without materializing the full score matrix.
+- `qkv_no_bias_materialized_sgmat_fused_f32` — Materializes FP32 scores for subgroup-matrix score and value products, then folds softmax normalization into the apply pass using row statistics emitted by the score pass.
+- `qkv_bias_materialized_sgmat_fused_f32` — Materializes FP32 scores for subgroup-matrix score and value products, then folds softmax normalization into the apply pass using row statistics emitted by the score pass.
+- `qkv_no_bias_materialized_sgmat_fused_f16` — Materializes FP32 scores for subgroup-matrix score and value products, then folds softmax normalization into the apply pass using row statistics emitted by the score pass.
+- `qkv_no_bias_flash_cluster_lpq4_nosg` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_no_bias_flash_cluster_nosg` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_bias_flash_cluster_nosg` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_no_bias_flash_cluster_lpq4` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_no_bias_flash_cluster` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_bias_flash_cluster` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_no_bias_flash_cluster_attn_bias` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_bias_flash_cluster_attn_bias` — Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.
+- `qkv_no_bias_flash_splitk_nosg` — Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.
+- `qkv_bias_flash_splitk_nosg` — Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.
+- `qkv_no_bias_flash_splitk` — Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.
+- `qkv_bias_flash_splitk` — Partitions the key reduction across workgroups and merges partial softmax statistics and values. Used when decode or another short query has too few independent query/head workgroups; the merge reassociates the key reduction.
+- `qkv_no_bias_materialized_gemm_f32` — Materializes float32 scores and applies tiled score, softmax, and value passes without requiring subgroup-matrix support.
+- `qkv_bias_materialized_gemm_f32` — Materializes float32 scores and applies tiled score, softmax, and value passes without requiring subgroup-matrix support.
+- `qkv_no_bias_materialized_gemm_fused_f32` — Materializes FP32 scores while publishing each row's maximum, combines those row statistics, then normalizes and applies the scores in one pass. This route is eligible for score matrices above the configured size threshold.
+- `qkv_bias_materialized_gemm_fused_f32` — Materializes FP32 scores while publishing each row's maximum, combines those row statistics, then normalizes and applies the scores in one pass. This route is eligible for score matrices above the configured size threshold.
+- `qkv_no_bias_attn_bias` — Portable scalar online attention with bias and additive-mask support. It applies online softmax while accumulating context, without materializing a score tensor.
+- `qkv_bias_attn_bias` — Portable scalar online attention with bias and additive-mask support. It applies online softmax while accumulating context, without materializing a score tensor.
+- `qkv_no_bias` — Portable scalar online attention with bias and additive-mask support. It applies online softmax while accumulating context, without materializing a score tensor.
+- `qkv_bias` — Portable scalar online attention with bias and additive-mask support. It applies online softmax while accumulating context, without materializing a score tensor.
+- `qkv_no_bias_flash` — Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_bias_flash` — Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_no_bias_flash_nosg` — Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_bias_flash_nosg` — Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_no_bias_flash_attn_bias` — Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_bias_flash_attn_bias` — Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_no_bias_flash_attn_bias_nosg` — Computes one-pass online-softmax attention without score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_bias_flash_attn_bias_nosg` — Computes one-pass online-softmax attention with fused projection bias and no score scratch. Family suffixes select subgroup, portable, and additive-attention-bias forms without changing the algorithm.
+- `qkv_no_bias_materialized_sgmat_f32` — Materializes FP32 scores for subgroup-matrix score and value products, with normalization performed by a separate softmax pass.
+- `qkv_bias_materialized_sgmat_f32` — Materializes FP32 scores for subgroup-matrix score and value products, with normalization performed by a separate softmax pass.
 ## 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
 - [`attn-materialized-softmax-f32.wgsl.jinja`](build/webgpu/attn-materialized-softmax-f32.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`attn-small-head-parallel.wgsl.jinja`](build/webgpu/attn-small-head-parallel.wgsl.jinja)
+- [`mha-small-seq-blocked.wgsl.jinja`](build/webgpu/mha-small-seq-blocked.wgsl.jinja)
 - [`mha-small-seq.wgsl.jinja`](build/webgpu/mha-small-seq.wgsl.jinja)
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 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
@@ -105,7 +106,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
   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,17 @@
 {% 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
@@ -81,20 +81,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
@@ -105,31 +107,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);
   }
@@ -144,9 +139,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;

 {% if cacheSeqlens is not defined %}{% set cacheSeqlens = false %}{% endif %}
 {% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
 {% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
+{% set scale = scale | default("0.0") %}
+{% if hasRotary is not defined %}{% set hasRotary = false %}{% endif %}
+{% set qSeq = qSeq | default(0) %}
+{% set splitKWorkgroupSize = workgroupSizeSpec if workgroupSizeSpec is defined else tunables.WORKGROUP_SIZE %}
 {% if useSubgroups %}
 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
 // 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;

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;
   let keyBound = keyBoundV;
   // Start at the WG-aligned tile containing keyFloor (sub-floor lanes mask out).
@@ -279,7 +261,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 (keyAllowed) {
       let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
 {% if hasMask %}
@@ -287,17 +269,6 @@ fn main(
       // [q, k] masks set batch/head strides to 0).
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qi * params.maskSeqStride + kj;
 {% endif %}
-{% if hasMask and maskIsBool %}
-      if (attn_mask[maskIndex] == 0u) {
-        keyAllowed = false;
-      } else {
-        var acc: f32 = 0.0;
-        for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
-          acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
-        }
-        score = acc * scale;
-      }
-{% else %}
       var acc: f32 = 0.0;
       for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
         acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
@@ -305,7 +276,6 @@ fn main(
       score = acc * scale;
 {% if hasMask %}
       score = score + f32(attn_mask[maskIndex]);
-{% endif %}
 {% endif %}
       m = score;
       dPart = select(0.0, 1.0, keyAllowed);

 {% 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 = select(params.kvSeq, min(params.kvSeq, qi + 1u), params.isCausal != 0u);
   var keyFloor: u32 = 0u;
   let keyBound = keyBoundV;
   // Start at the WG-aligned tile containing keyFloor (sub-floor lanes mask out).
     var score = -FLT_MAX;
     var m = -FLT_MAX;
     var dPart = 0.0;
+    var keyAllowed = kj < keyBound && (params.isCausal == 0u || kj <= qi);
     if (keyAllowed) {
       let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
 {% if hasMask %}
       // [q, k] masks set batch/head strides to 0).
       let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qi * params.maskSeqStride + kj;
 {% endif %}
       var acc: f32 = 0.0;
       for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
         acc = acc + dot(q_shared[d4], vec4<f32>(key[kRowV4 + d4]));
       score = acc * scale;
 {% if hasMask %}
       score = score + f32(attn_mask[maskIndex]);
 {% endif %}
       m = score;
       dPart = select(0.0, 1.0, keyAllowed);

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 %}
@@ -27,9 +30,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
@@ -43,12 +43,14 @@ 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) %}
 // 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
@@ -62,7 +64,8 @@ enable f16;
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
-const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const TILE_Q: u32 = {{ TILE_Q }}u;
@@ -70,7 +73,12 @@ const TILE_Q: u32 = {{ TILE_Q }}u;
 const LPQ: u32 = {{ LPQ }}u;
 const SLICE: u32 = HEAD_DIM / (4u * LPQ);  // vec4s per lane slice
 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
@@ -103,8 +111,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 %}
@@ -171,36 +183,48 @@ 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);
-  let qBase4 = (b * {{ QSEQ }} + qClamped) * {{ Q_STRIDE }} + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% 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 %}
-  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.
-{% if hasCausal or hasWindow %}
   // Upper-left causal and/or sliding-window bounds. Query qIdx sits at
   // absolute position p = kvSeq - qSeq + qIdx, so it attends keys [minKj, maxKj):
   // the causal ceiling p+1 (clamped to kvSeq for non-square chunks) and, when
@@ -209,36 +233,26 @@ fn main(
   // pastLen = kvSeq - qSeq: query tokens are the last qSeq positions of the cache;
   // saturate so an out-of-contract qSeq > kvSeq can't underflow u32. Used for the
   // sliding window always, and for the causal ceiling only when right-aligned.
-{% if hasCausal %}
-  var maxKjV = select({{ KVA }}, min(qIdx + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
-{% else %}
-  var maxKjV = {{ KVA }};
-{% endif %}
-  var minKjV: u32 = 0u;
   let lastQ = min(wg.x * TILE_Q + TILE_Q - 1u, {{ QSEQ }} - 1u);
-{% if hasCausal %}
   var wgEndV = select({{ KVA }}, min(lastQ + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
-{% else %}
-  var wgEndV = {{ KVA }};
-{% endif %}
   var wgStartV: u32 = 0u;
-  let maxKj = maxKjV;
-  let minKj = minKjV;
   let wgEnd = wgEndV;
   let wgStart = wgStartV;
-{% else %}
-  let maxKj = {{ KVA }};
-  let minKj: u32 = 0u;
-  let wgEnd = {{ KVA }};
-  let wgStart: u32 = 0u;
-{% endif %}
   let kvBatch4 = b * {{ KVSEQ }} * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   var kStart: u32 = wgStart;
   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;
@@ -274,8 +288,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>;
@@ -330,6 +343,67 @@ fn main(
 {% endif %}
       }
     }
 {% else %}
     var s: array<f32, TILE_K>;
     for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
@@ -392,30 +466,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) %}
     {
@@ -432,31 +531,34 @@ fn main(
       o{{ c }} = acc;
     }
 {% endfor %}
     kStart = kStart + TILE_K;
   }
-  if (qValid) {
-    let outBase4 = (b * {{ QSEQ }} + qIdx) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% if hasHeadSink %}
     // The head sink is a learned logit that competes with the keys but carries
     // no value, so it enters the denominator only and the weighted sum above is
-    // untouched. Renormalizing against max(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 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 %}
 // 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: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
 const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
+{% if sourceProfile == 1 %}const QKV_STRIDE_V4: u32 = {{ qkvStrideV4 }}u; // packed [Q; K; V] input row stride
+{% endif %}const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
 const TILE_Q: u32 = {{ TILE_Q }}u;
 const LPQ: u32 = {{ LPQ }}u;
 const SLICE: u32 = HEAD_DIM / (4u * LPQ);  // vec4s per lane slice
 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 %}
   // 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);
+  let {{ qn("qBase4", qi) }} = (b * {{ QSEQ }} + {{ qn("qClamped", qi) }}) * {{ Q_STRIDE }} + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% for c in range(SLICE_COUNT) %}
+  var {{ qn("qr" ~ c, qi) }} = vec4<f32>({{ QUERY }}[{{ qn("qBase4", qi) }} + {{ c }}u]);
 {% if hasBias %}
+  {{ qn("qr" ~ c, qi) }} = {{ qn("qr" ~ c, qi) }} + load_bias4(h * HEAD_DIM, lane8 * SLICE + {{ c }}u);
 {% endif %}
+  var {{ qn("o" ~ c, qi) }} = vec4<f32>(0.0);
+{% endfor %}
+{% endfor %}
+{% 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.
   // Upper-left causal and/or sliding-window bounds. Query qIdx sits at
   // absolute position p = kvSeq - qSeq + qIdx, so it attends keys [minKj, maxKj):
   // the causal ceiling p+1 (clamped to kvSeq for non-square chunks) and, when
   // pastLen = kvSeq - qSeq: query tokens are the last qSeq positions of the cache;
   // saturate so an out-of-contract qSeq > kvSeq can't underflow u32. Used for the
   // sliding window always, and for the causal ceiling only when right-aligned.
+{% for qi in range(QPL) %}
+  var {{ qn("maxKjV", qi) }} = select({{ KVA }}, min({{ qn("qIdx", qi) }} + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
+  var {{ qn("minKjV", qi) }}: u32 = 0u;
+{% endfor %}
   let lastQ = min(wg.x * TILE_Q + TILE_Q - 1u, {{ QSEQ }} - 1u);
   var wgEndV = select({{ KVA }}, min(lastQ + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
   var wgStartV: u32 = 0u;
+{% 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;
   let kvBatch4 = b * {{ KVSEQ }} * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
   var kStart: u32 = wgStart;
   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{% if hasBias %} + load_bias4(2u * {{ Q_HIDDEN }}u + h * HEAD_DIM, lane8 * SLICE + {{ c }}u){% endif %}{% endmacro %}
+{% for qi in range(QPL) %}
+  if ({{ qn("qValid", qi) }}) {
+    let outBase4 = (b * {{ QSEQ }} + {{ qn("qIdx", qi) }}) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
 {% if hasHeadSink %}
     // The head sink is a learned logit that competes with the keys but carries
     // no value, so it enters the denominator only and the weighted sum above is
+    // untouched. Renormalizing against max({{ qn("m", qi) }}, sink) keeps the exponentials in
     // range when the sink dominates a fully-masked row.
     let sink = f32(head_sink[h]);
+    let finalM = max({{ qn("m", qi) }}, sink);
+    let accScale = exp({{ qn("m", qi) }} - finalM);
+    let inv = accScale / (exp(sink - finalM) + {{ qn("l", qi) }} * accScale);
 {% else %}
+    // {{ qn("l", qi) }} == 0 means this query had no probability-bearing key: either its
     // causal/window range is empty or its logical bool mask rejects every key.
     // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
     // contract by selecting on positive global mass.
+    let inv = select(0.0, 1.0 / {{ qn("l", qi) }}, {{ qn("l", qi) }} > 0.0);
 {% endif %}
 {% for c in range(SLICE_COUNT) %}
+    {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c, qi) }});
 {% endfor %}
   }
+{% endfor %}
 }

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

@@ -9,12 +9,9 @@
 // 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 %}

 // 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 %}
 {% set CAUSAL = false if (hasCausal is defined and not hasCausal) else true %}
 {% if USE_SUBGROUPS %}
 enable subgroups;
 {% endif %}

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

@@ -1,16 +1,14 @@
 {{ env.wgsl.resourceDeclarations }}
-// Register-blocked normalized-scores @ V GEMM for BHSD Attention and BSH MHA.
 // Score rows are head-major regardless of the external tensor layout.
 //
 // V is stored row-major in the head dimension, so its tile groups four of those
 // per vector word and the micro-tile accumulates four output columns at a time.
-// The score tile groups four KEYS per word -- the axis this GEMM reduces over --
-// so one step over four keys reads TM + TN vector words for TM * TN * 4 fused
-// multiply-adds. Grouping the reduction axis is what keeps the read-per-FMA
-// ratio at the score GEMM's: a key-scalar tile costs three times as many
-// workgroup reads for the same arithmetic.
-{% set FUSED_SOFTMAX = source.fusedSoftmax is defined and source.fusedSoftmax %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;

 {{ env.wgsl.resourceDeclarations }}
+// Register-blocked normalized-scores @ V GEMM with BHSD and BSH specializations.
 // Score rows are head-major regardless of the external tensor layout.
 //
 // V is stored row-major in the head dimension, so its tile groups four of those
 // per vector word and the micro-tile accumulates four output columns at a time.
+// The score tile groups four keys per word along the reduction axis. One step
+// over four keys reads TM + TN vector words for TM * TN * 4 fused
+// multiply-adds.
+{% set FUSED_SOFTMAX = fusedSoftmax is defined and fusedSoftmax %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;

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

@@ -12,7 +12,7 @@
 // `maxOnly` means the producer published only a row max per slot because
 // computing the denominator there would double the exp count. Fold maxima and
 // leave the denominator to the apply pass, which sees every row element anyway.
-{% set MAX_ONLY = source.maxOnly is defined and source.maxOnly %}
 {{ env.wgsl.resourceDeclarations }}
 const SLOTS: u32 = {{ statSlots }}u;
@@ -52,10 +52,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

 // `maxOnly` means the producer published only a row max per slot because
 // computing the denominator there would double the exp count. Fold maxima and
 // leave the denominator to the apply pass, which sees every row element anyway.
+{% set MAX_ONLY = maxOnly is defined and maxOnly %}
 {{ env.wgsl.resourceDeclarations }}
 const SLOTS: u32 = {{ statSlots }}u;
 @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-score-f32.wgsl.jinja CHANGED Viewed

@@ -1,21 +1,20 @@
 {{ env.wgsl.resourceDeclarations }}
-// f32 prefill score GEMM for rank-4 Attention (BHSD) and packed MHA (BSH).
-// Materializing the score matrix is reserved for large, compute-bound prefills:
-// it turns the shuffle/reduction-heavy online QK loop into a regular 64x64
-// register-blocked GEMM. Q and K are both stored row-major in the head
 // dimension — the axis this GEMM reduces over — so both tiles group four of
 // its values per vector word, every staging read is contiguous, and the
 // micro-tile accumulates through dot(): one step reads TM + TN vector words
 // for TM * TN * 4 fused multiply-adds. K keeps a one-word shared stride
-// padding, which avoids the worst transpose-read bank conflicts.
 {% if scalarAccumulators is not defined %}{% set scalarAccumulators = false %}{% endif %}
 {% if maskIsKeyKeep is not defined %}{% set maskIsKeyKeep = false %}{% endif %}
 {% if scalarAccumulators %}
 {% set TM_JINJA = (materializedQueryTile / materializedWorkgroupDim)|int %}
 {% set TN_JINJA = (materializedKeyTile / materializedWorkgroupDim)|int %}
 {% endif %}
-{% set EMIT_ROW_STATS = source.emitRowStats is defined and source.emitRowStats %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
@@ -109,8 +108,7 @@ fn main(
   let li = lid.y * WG_DIM + lid.x;
 {% if scalarAccumulators %}
-  // Keep every register-tile cell statically addressable. A dynamically indexed
-  // private array may fail to scalarize and spill hot GEMM accumulators to memory.
 {% for i in range(TM_JINJA) %}
 {% for j in range(TN_JINJA) %}
   var acc{{ i }}_{{ j }}: f32 = 0.0;

 {{ env.wgsl.resourceDeclarations }}
+// f32 prefill score GEMM with BHSD and packed BSH specializations. It computes
+// the score matrix as a 64x64 register-blocked product. Q and K are both stored row-major in the head
 // dimension — the axis this GEMM reduces over — so both tiles group four of
 // its values per vector word, every staging read is contiguous, and the
 // micro-tile accumulates through dot(): one step reads TM + TN vector words
 // for TM * TN * 4 fused multiply-adds. K keeps a one-word shared stride
+// padding, which changes the transposed workgroup-memory stride and reduces
+// bank-conflict risk on banked implementations.
 {% if scalarAccumulators is not defined %}{% set scalarAccumulators = false %}{% endif %}
 {% if maskIsKeyKeep is not defined %}{% set maskIsKeyKeep = false %}{% endif %}
 {% if scalarAccumulators %}
 {% set TM_JINJA = (materializedQueryTile / materializedWorkgroupDim)|int %}
 {% set TN_JINJA = (materializedKeyTile / materializedWorkgroupDim)|int %}
 {% endif %}
+{% set EMIT_ROW_STATS = emitRowStats is defined and emitRowStats %}
 const Q_HEADS: u32 = {{ qNumHeads }}u;
 const KV_HEADS: u32 = {{ kvNumHeads }}u;
   let li = lid.y * WG_DIM + lid.x;
 {% if scalarAccumulators %}
+  // Keep every register-tile cell statically addressable.
 {% for i in range(TM_JINJA) %}
 {% for j in range(TN_JINJA) %}
   var acc{{ i }}_{{ j }}: f32 = 0.0;

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

@@ -9,17 +9,18 @@ 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" %}
-{% macro q_index(row, d) %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endmacro %}
-{% macro kv_index(seq, d) %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% 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 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 %}
@@ -44,7 +45,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 %}
@@ -63,7 +64,9 @@ const KV_HEADS: u32 = {{ kvNumHeads }}u;
 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
@@ -71,9 +74,11 @@ const HIDDEN: u32 = {{ qHidden }}u;
  * constant. V is token-independent, so sum_k p[k] * bv = bv after
  * normalization; it is added after the apply GEMM. */
 {% endif %}
 // Packed K/V rows span KV_HEADS heads, so they are narrower than the query
 // row whenever queries are grouped. Q and the output keep HIDDEN.
 const KV_HIDDEN: u32 = KV_HEADS * HEAD_DIM;
 {% if EMIT_ROW_STATS %}
 const STAT_SLOTS: u32 = {{ statSlots }}u;
 {% endif %}
@@ -191,7 +196,7 @@ fn main(
 {% endif %}
 {% endif %}
   for (var k_base = 0u; k_base < inner; 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 (
@@ -229,10 +234,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++) {
@@ -304,10 +307,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;

 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" %}
+{% 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 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 %}
 const KV_HEADS: u32 = HEADS;
 {% endif %}
 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
  * constant. V is token-independent, so sum_k p[k] * bv = bv after
  * normalization; it is added after the apply GEMM. */
 {% endif %}
+{% if not kvHeadMajor %}
 // Packed K/V rows span KV_HEADS heads, so they are narrower than the query
 // row whenever queries are grouped. Q and the output keep HIDDEN.
 const KV_HIDDEN: u32 = KV_HEADS * HEAD_DIM;
+{% endif %}
 {% if EMIT_ROW_STATS %}
 const STAT_SLOTS: u32 = {{ statSlots }}u;
 {% endif %}
 {% endif %}
 {% endif %}
   for (var k_base = 0u; k_base < inner; k_base += TILE_K) {
+{% if phase == "apply" and not FUSED_SOFTMAX and MT == "f32" %}
     // Full interior PV tiles can be loaded directly from storage. Query,
     // reduction, and output-dimension tails use the guarded shared path below.
     if (
 {% 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;

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

@@ -1,4 +1,4 @@
-{% set CACHE_VEC4 = source.cacheVec4 if source.cacheVec4 is defined else false %}
 {% set CAUSAL_ROWS = causalRows is defined and causalRows %}
 {% set SCALE_IN = scaleInSoftmax is defined and scaleInSoftmax %}
 {% set RUNTIME_COLS = colsFromParams is defined and colsFromParams %}
@@ -50,20 +50,22 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 // 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
@@ -74,31 +76,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);
   }
@@ -113,9 +108,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;
@@ -177,12 +171,11 @@ fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
   @builtin(subgroup_size) sgSize: u32
 {% endif %}
 ) {
-  let row = wg.x + wg.y * nwg.x;
   if (row >= {{ ROWS_EXPR }}) { return; }
 {% if CACHE_VEC4 %}
   let base4 = row * COLS4;

+{% set CACHE_VEC4 = cacheVec4 if cacheVec4 is defined else false %}
 {% set CAUSAL_ROWS = causalRows is defined and causalRows %}
 {% set SCALE_IN = scaleInSoftmax is defined and scaleInSoftmax %}
 {% set RUNTIME_COLS = colsFromParams is defined and colsFromParams %}
 // 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;
 @compute @workgroup_size(WG, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
   @builtin(subgroup_size) sgSize: u32
 {% endif %}
 ) {
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= {{ ROWS_EXPR }}) { return; }
 {% if CACHE_VEC4 %}
   let base4 = row * COLS4;

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

@@ -1,25 +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 %}
@@ -30,10 +19,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 +67,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 +113,7 @@ 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 %}
   for (var key_token: u32 = minKj; key_token < maxKj; key_token = key_token + 1u) {
     let kRow = kvBase + key_token * kvTokenStride;
@@ -183,7 +159,9 @@ fn main(
     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 hasBias %}
 // Packed [Q; K; V] bias rows are applied during the serial key walk.
 {% endif %}
 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;
   maxKj = min(maxKj, select(params.kvSeq, query_token + 1u, params.isCausal != 0u));
   for (var key_token: u32 = minKj; key_token < maxKj; key_token = key_token + 1u) {
     let kRow = kvBase + key_token * kvTokenStride;
     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/attn-small-head-parallel.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Small-head attention is dominated by the per-key barriers in the generic





1	{{ env.wgsl.resourceDeclarations }}
2
3	// Small-head attention is dominated by the per-key barriers in the generic

build/webgpu/bench.json CHANGED Viewed

@@ -1,10 +1,9 @@
 {
-  "op": "com.microsoft.MultiHeadAttention",
   "cases": [
     {
       "name": "minilm-self-attention-b1-s20-h12-d32",
       "provenance": {
-        "notes": "Upper boundary of the whole-head small-sequence route at 12 heads and sequence length 20. Paired with the sequence-24 case, it checks that the mutually exclusive route predicates meet without a gap."
       },
       "vars": { "batch": 1, "qSeq": 20, "kvSeq": 20, "heads": 12, "headDim": 32 },
       "attrs": { "num_heads": 12 },

 {
   "cases": [
     {
       "name": "minilm-self-attention-b1-s20-h12-d32",
       "provenance": {
+        "notes": "Twelve heads at sequence length 20 mark the upper boundary of the whole-head small-sequence route; `minilm-self-attention-b1-s24-h12-d32` exercises the adjacent flash route."
       },
       "vars": { "batch": 1, "qSeq": 20, "kvSeq": 20, "heads": 12, "headDim": 32 },
       "attrs": { "num_heads": 12 },

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,31 +1,78 @@
 {
   "name": "com.microsoft.MultiHeadAttention",
-  "id": "_com_microsoft_multiheadattention_webgpu_03bfaa6",
   "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": "8RI/yP6118n/rLRo0OG9NJ+kZ2xdmJyj093Tw8C07nc=",
-      "attn-flash-decode-splitk.wgsl.jinja": "Pyt7ZMMGU2H8ap4T4ggVRAhW38fotpNyJKr7z+iVsBU=",
-      "attn-flash-online.wgsl.jinja": "uJ68uHbQN9k8TLdgtZtdkw78iXYgTnE1UpWdt1LGq7w=",
-      "attn-flash-prefill-cluster.wgsl.jinja": "TO+SM9D9yK/TC4LOL/3U4/xHd0HS6xB61S+POu9p2QQ=",
-      "attn-flash-q32-broadcast.wgsl.jinja": "evSdaiYgvrikYZ8hpb5m4gxf+rIryKkHOTiKgyF7ioI=",
-      "attn-materialized-apply-f32.wgsl.jinja": "qbrsh9S14QYuAM+xF49dgcooO6f19yNg9lC8eoBeWFQ=",
-      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "rdpdnh6y3m0tBwFMVCc903j2cvGsnGTcT8EAWRr5TRk=",
-      "attn-materialized-score-f32.wgsl.jinja": "EAhU/wHSgmJZ/YHCnwG/takpgfMAtGqKXo3Cf8TfvxU=",
-      "attn-materialized-sgmat-f32.wgsl.jinja": "7fJFDK+kWEhN4rqFs3fE9uPkZdHh5C9xl7PXe9vZze4=",
-      "attn-materialized-softmax-f32.wgsl.jinja": "7z288syeGhUUyBpaZzVVqdtFVMQALuN88sNfytvClWw=",
-      "attn-online-scalar.wgsl.jinja": "CW0ru5yg8YFpIv31XnnD5IuKBXppsKc+2BU7u2uAQv4=",
-      "attn-small-head-parallel.wgsl.jinja": "ejotx9NjdqElO3EIQWj8E3hJk2iW8Exq1vUjxK4FUk0=",
-      "bench.json": "VV4kVBFJ112r3U2RGQlQd1Z5eAzgm+y30+oGG6TpK3g=",
-      "manifest.json": "kaCZGh3/PB8EmCgBc5QCc0aVGSrjf55xB6OM4s0yhgg=",
-      "mha-small-seq.wgsl.jinja": "F825nrB6RLu55ziDnDI7BTJ3JSdT1GwsuMYmtrHQSic=",
-      "test.json": "6PCZYazz1hXA9+nbVJmd0/btWsIkkp29/xYqWZ4OSrs="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MultiHeadAttention" }
 }

 {
   "name": "com.microsoft.MultiHeadAttention",
+  "id": "_com_microsoft_multiheadattention_webgpu_5144476",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "attention-rank4-tiled.wgsl.jinja": "gYl58ecIVTxJJkvnK/6lNcbPKaRT24cfz5Yx0MxxxxE=",
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "XSAMF/uzQKJg9pZZDeqmViaVXuj47Vd7Pk404Z1J48I=",
+      "attn-flash-decode-splitk.wgsl.jinja": "dqxX8taw1tpo/QQkFr2zJ7TEEtQtOspTYEYHADXtESs=",
+      "attn-flash-online.wgsl.jinja": "hCxO/6VGkyixwyoT/kuYeHgcEkFoRwk/ZmiLlnwa0jM=",
+      "attn-flash-prefill-cluster.wgsl.jinja": "igbsg5xRuFz7alnwmL74hhh1JUqfDB5PfUslsEz5rjA=",
+      "attn-flash-q32-broadcast.wgsl.jinja": "8qgKXydwiH2shac/fkdHi+FR+c8pOMDFbNp9Zs8KAnc=",
+      "attn-materialized-apply-f32.wgsl.jinja": "aF7JSJj8rQN8wGTV7IWmxkqMDbXOnGBSI3tsI1CxZNk=",
+      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "+/5m03sN44zTn9q7cquDea2UNhEO1ujzm1J7yPai4OY=",
+      "attn-materialized-score-f32.wgsl.jinja": "gCvBC1T1CLEmNxcpo04CLBQdFBkEVMJ9Xhv1wbQxCsI=",
+      "attn-materialized-sgmat-f32.wgsl.jinja": "Ez9UJggPXMZ6OyzUQu8xLtg8bZd+lTCoEGDwLkNHNL0=",
+      "attn-materialized-softmax-f32.wgsl.jinja": "DCuz4gFKXDCdR9UMnfD+mYFjBghnv12GxtEsPKH52mE=",
+      "attn-online-scalar.wgsl.jinja": "6+hU7MzwHz69fjciv94xsPxN+o9rwm42u9Cin2C+AL8=",
+      "attn-small-head-parallel.wgsl.jinja": "oxFE6hQZKTRugz9hwuS+cob00PskT8iYqP4e+YWtHKA=",
+      "bench.json": "b/WdaAhEP5NpRXNB5w1mivRNTfALrBDJisDONxTf4Yo=",
+      "manifest.json": "+/bNCAjkG5OoutFq7O4vQJqb2yMxzICsKvFXsRE50bA=",
+      "mha-small-seq-blocked.wgsl.jinja": "zznGERrPco5uH9tKiQ2QE+ArO3vho/nEv69wT8IgY1U=",
+      "mha-small-seq.wgsl.jinja": "UsbUIkvgWL/RgOBSm7lhPGIatA36yTQR4Z2gpHVLqx4=",
+      "test.json": "NlMA/wlmGSpogAuhSZ/kcQ01dd3M+CvkPF0d/VBmIpM="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "qkv_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
+      "qkv_no_bias_small_seq_blocked": ["mha-small-seq-blocked.wgsl.jinja"],
+      "qkv_no_bias_small_seq": ["mha-small-seq.wgsl.jinja"],
+      "qkv_no_bias_tiled_nosg": ["attention-rank4-tiled.wgsl.jinja"],
+      "qkv_bias_flash_q32_broadcast_f32_d128": ["attn-flash-q32-broadcast.wgsl.jinja"],
+      "qkv_no_bias_small_head_parallel": ["attn-small-head-parallel.wgsl.jinja"],
+      "qkv_no_bias_tiled_attn_bias_nosg": ["attention-rank4-tiled.wgsl.jinja"],
+      "qkv_no_bias_materialized_sgmat_fused_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja"],
+      "qkv_bias_materialized_sgmat_fused_f32": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja"],
+      "qkv_no_bias_materialized_sgmat_fused_f16": ["attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-sgmat-f32.wgsl.jinja"],
+      "qkv_no_bias_flash_cluster_lpq4_nosg": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_no_bias_flash_cluster_nosg": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_bias_flash_cluster_nosg": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_no_bias_flash_cluster_lpq4": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_no_bias_flash_cluster": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_bias_flash_cluster": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_no_bias_flash_cluster_attn_bias": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_bias_flash_cluster_attn_bias": ["attn-flash-prefill-cluster.wgsl.jinja"],
+      "qkv_no_bias_flash_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "qkv_bias_flash_splitk_nosg": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "qkv_no_bias_flash_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "qkv_bias_flash_splitk": ["attn-flash-decode-splitk-merge.wgsl.jinja", "attn-flash-decode-splitk.wgsl.jinja"],
+      "qkv_no_bias_materialized_gemm_f32": ["attn-materialized-apply-f32.wgsl.jinja", "attn-materialized-score-f32.wgsl.jinja", "attn-materialized-softmax-f32.wgsl.jinja"],
+      "qkv_bias_materialized_gemm_f32": ["attn-materialized-apply-f32.wgsl.jinja", "attn-materialized-score-f32.wgsl.jinja", "attn-materialized-softmax-f32.wgsl.jinja"],
+      "qkv_no_bias_materialized_gemm_fused_f32": ["attn-materialized-apply-f32.wgsl.jinja", "attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-score-f32.wgsl.jinja"],
+      "qkv_bias_materialized_gemm_fused_f32": ["attn-materialized-apply-f32.wgsl.jinja", "attn-materialized-rowstats-combine-f32.wgsl.jinja", "attn-materialized-score-f32.wgsl.jinja"],
+      "qkv_no_bias_flash_q32_broadcast": ["attn-flash-q32-broadcast.wgsl.jinja"],
+      "qkv_no_bias_flash_q32_shared": ["attn-flash-q32-broadcast.wgsl.jinja"],
+      "qkv_no_bias_attn_bias": ["attn-online-scalar.wgsl.jinja"],
+      "qkv_bias_attn_bias": ["attn-online-scalar.wgsl.jinja"],
+      "qkv_no_bias": ["attn-online-scalar.wgsl.jinja"],
+      "qkv_bias": ["attn-online-scalar.wgsl.jinja"],
+      "qkv_no_bias_flash": ["attn-flash-online.wgsl.jinja"],
+      "qkv_bias_flash": ["attn-flash-online.wgsl.jinja"],
+      "qkv_no_bias_flash_nosg": ["attn-flash-online.wgsl.jinja"],
+      "qkv_bias_flash_nosg": ["attn-flash-online.wgsl.jinja"],
+      "qkv_no_bias_flash_attn_bias": ["attn-flash-online.wgsl.jinja"],
+      "qkv_bias_flash_attn_bias": ["attn-flash-online.wgsl.jinja"],
+      "qkv_no_bias_flash_attn_bias_nosg": ["attn-flash-online.wgsl.jinja"],
+      "qkv_bias_flash_attn_bias_nosg": ["attn-flash-online.wgsl.jinja"],
+      "qkv_no_bias_materialized_sgmat_f32": ["attn-materialized-sgmat-f32.wgsl.jinja", "attn-materialized-softmax-f32.wgsl.jinja"],
+      "qkv_bias_materialized_sgmat_f32": ["attn-materialized-sgmat-f32.wgsl.jinja", "attn-materialized-softmax-f32.wgsl.jinja"]
+    }
+  }
 }

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

	@@ -0,0 +1,166 @@

+{{ env.wgsl.resourceDeclarations }}
+// Blocked whole-head attention for short key sequences with optional projection
+// bias and causal masking. One workgroup owns
+// (batch, head) and a block of QUERY_BLOCK query rows; it stages the head's
+// complete K and V planes in workgroup memory once, then gives every query
+// KEY_LANES lanes: lane (q, kl) scores keys kl, kl + KEY_LANES, ... against
+// its private Q row, keeps an online-softmax (max, denominator, context)
+// partial in registers, and the query's KEY_LANES partials merge through
+// workgroup memory in one pass. The full score row never exists.
+//
+// The complete (query, key) block is folded with one reduction round per query
+// block. KEY_LANES independently accumulate interleaved keys before that fold.
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const HIDDEN_V4: u32 = {{ hiddenV4 }}u;
+const KV_SEQ: u32 = {{ kvSeq }}u;
+const QUERY_BLOCK: u32 = {{ queryBlock }}u;
+const KEY_LANES: u32 = {{ keyLanes }}u;
+const WG: u32 = QUERY_BLOCK * KEY_LANES;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% if hasBias %}
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
+}
+{% endif %}
+var<workgroup> kShared: array<vec4<f32>, KV_SEQ * HEAD_DIM_V4>;
+var<workgroup> vShared: array<vec4<f32>, KV_SEQ * HEAD_DIM_V4>;
+// Per-lane partials for the merge: (max, denom) and the context row.
+var<workgroup> partM: array<f32, WG>;
+var<workgroup> partD: array<f32, WG>;
+var<workgroup> partO: array<vec4<f32>, WG * HEAD_DIM_V4>;
+@compute @workgroup_size({{ queryBlock * keyLanes }}, 1, 1)
+fn main(
+  @builtin(workgroup_id) wid: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let head = wid.y;
+  let batch = wid.z;
+  let tid = lid.x;
+  let ql = tid / KEY_LANES;
+  let kl = tid % KEY_LANES;
+  let qi = wid.x * QUERY_BLOCK + ql;
+  let headV4 = head * HEAD_DIM_V4;
+  // Lanes stage this head's K and V planes from consecutive vec4 elements,
+  // advancing by the workgroup width even when their query row is absent.
+  let kvBase = batch * KV_SEQ * HIDDEN_V4;
+  for (var i = tid; i < KV_SEQ * HEAD_DIM_V4; i = i + WG) {
+    let s = i / HEAD_DIM_V4;
+    let d4 = i % HEAD_DIM_V4;
+    let src = kvBase + s * HIDDEN_V4 + headV4 + d4;
+    kShared[i] = vec4<f32>(key[src]);
+    vShared[i] = vec4<f32>(value[src]);
+  }
+  workgroupBarrier();
+  let scale = scale_value();
+  let rowValid = qi < params.qSeq;
+  // A query past the end scores no keys: its partials are the identity and it
+  // stores nothing, so every barrier below stays uniform. For a causal request,
+  // query qi attends keys 0 through qi.
+  let keyBound = select(0u, select(KV_SEQ, min(KV_SEQ, qi + 1u), params.isCausal != 0u), rowValid);
+  // Private Q row (with the Q bias folded in), read from the clamped row so an
+  // absent query still reads in bounds.
+  let qRowV4 = (batch * params.qSeq + min(qi, params.qSeq - 1u)) * HIDDEN_V4 + headV4;
+{% for d in range(headDimV4) %}
+  var q{{ d }} = vec4<f32>(query[qRowV4 + {{ d }}u]);
+{% if hasBias %}
+  q{{ d }} = q{{ d }} + load_bias4(head * HEAD_DIM, {{ d }}u);
+{% endif %}
+{% endfor %}
+  // Lane-local online softmax over keys kl, kl + KEY_LANES, ...
+  var m = -FLT_MAX;
+  var d = 0.0;
+{% for c in range(headDimV4) %}
+  var o{{ c }} = vec4<f32>(0.0);
+{% endfor %}
+  for (var kj = kl; kj < keyBound; kj = kj + KEY_LANES) {
+    let kBase = kj * HEAD_DIM_V4;
+    var s = 0.0;
+{% for c in range(headDimV4) %}
+    s = s + dot(q{{ c }}, kShared[kBase + {{ c }}u]);
+{% endfor %}
+    s = s * scale;
+    let mNew = max(m, s);
+    let corr = exp_shift(m, mNew);
+    let p = exp_shift(s, mNew);
+    d = d * corr + p;
+{% for c in range(headDimV4) %}
+    o{{ c }} = o{{ c }} * corr + p * vShared[kBase + {{ c }}u];
+{% endfor %}
+    m = mNew;
+  }
+  partM[tid] = m;
+  partD[tid] = d;
+{% for c in range(headDimV4) %}
+  partO[tid * HEAD_DIM_V4 + {{ c }}u] = o{{ c }};
+{% endfor %}
+  workgroupBarrier();
+  // The query's leader lane merges its KEY_LANES partials with the stable
+  // online rule, then applies the V bias once and stores the row.
+  if (kl == 0u && rowValid) {
+    let base = ql * KEY_LANES;
+    var mAll = -FLT_MAX;
+    for (var i = 0u; i < KEY_LANES; i = i + 1u) {
+      mAll = max(mAll, partM[base + i]);
+    }
+    var dAll = 0.0;
+{% for c in range(headDimV4) %}
+    var out{{ c }} = vec4<f32>(0.0);
+{% endfor %}
+    for (var i = 0u; i < KEY_LANES; i = i + 1u) {
+      let w = exp_shift(partM[base + i], mAll);
+      dAll = dAll + partD[base + i] * w;
+{% for c in range(headDimV4) %}
+      out{{ c }} = out{{ c }} + partO[(base + i) * HEAD_DIM_V4 + {{ c }}u] * w;
+{% endfor %}
+    }
+    let hasMass = dAll > 0.0;
+    let inv = select(0.0, 1.0 / dAll, hasMass);
+    let outRowV4 = (batch * params.qSeq + qi) * HIDDEN_V4 + headV4;
+{% for c in range(headDimV4) %}
+    var v{{ c }} = out{{ c }} * inv;
+{% if hasBias %}
+    // V bias is token-independent, so add it once after the softmax-weighted sum.
+    v{{ c }} = v{{ c }} + load_bias4(2u * {{ hidden }}u + head * HEAD_DIM, {{ c }}u) * select(0.0, 1.0, hasMass);
+{% endif %}
+    output[outRowV4 + {{ c }}u] = v{{ c }};
+{% endfor %}
+  }
+}

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

@@ -28,9 +28,9 @@ fn main(
   let batch = wid.y;
   let headBase = head * HEAD_DIM;
-  // Cooperative, coalesced load of this head's K and V planes. Every thread
-  // participates regardless of whether it owns a query row, and the barrier is
-  // reached by the whole workgroup before any row exits.
   for (var i = lid.x; i < KV_SEQ * HEAD_DIM; i += WG) {
     let s = i / HEAD_DIM;
     let d = i % HEAD_DIM;

   let batch = wid.y;
   let headBase = head * HEAD_DIM;
+  // Lanes stage this head's K and V planes from consecutive flattened elements,
+  // advancing by the workgroup width. Every lane reaches the barrier before a
+  // lane without a query row exits.
   for (var i = lid.x; i < KV_SEQ * HEAD_DIM; i += WG) {
     let s = i / HEAD_DIM;
     let d = i % HEAD_DIM;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.MultiHeadAttention",
   "cases": [
     {
       "name": "cross_attention_no_bias",
@@ -186,7 +185,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/multihead_attention_op_test.cc",
         "test": "MultiHeadAttentionTest.CrossAttention_Batch1_HeadSize16",
-        "notes": "No-bias companion for ORT's batch=1, q_seq=2, kv_seq=3, num_heads=2, head_size=16 cross-attention shape."
       },
       "attrs": { "num_heads": 2 },
       "inputs": {
@@ -279,7 +278,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 64, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 64, 256], "tolerance": 0.0002 } }
@@ -350,7 +349,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0002 } }
@@ -399,7 +398,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0003 } }
@@ -407,7 +406,7 @@
     {
       "name": "splitk_decode_head_dim128_no_bias_kv2048",
       "provenance": {
-        "notes": "Compact companion for bench case decode-h8-d128-kv2048, covering headDim=128 decode split-K selection."
       },
       "attrs": { "num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
@@ -424,7 +423,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0003 } }
@@ -446,7 +445,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [2, 768, 512],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [2, 1, 512], "tolerance": 0.0003 } }
@@ -508,7 +507,7 @@
     {
       "name": "flash_self_attention_head_dim128_no_bias_f16",
       "provenance": {
-        "notes": "f16 headDim=128 prefill coverage for q32-broadcast benchmarks; full tier selects qkv_no_bias_flash_q32_broadcast and no-subgroups selects qkv_no_bias_flash_q32_shared."
       },
       "attrs": { "num_heads": 4, "scale": 0.08838834764831845 },
       "inputs": {
@@ -525,10 +524,10 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 64, 512],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
-      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 64, 512], "tolerance": 0.03 } }
     },
     {
       "name": "flash_decode_head_dim64_no_bias_f16",
@@ -547,10 +546,10 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 512, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
-      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 256], "tolerance": 0.03 } }
     },
     {
       "name": "cross_attention_projection_bias_f16",
@@ -583,7 +582,7 @@
       "name": "causal_cross_attention_q2_kv4_exact",
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/attention_helper.h",
-        "notes": "Exact upper-left causal witness. ORT CPU skips masking the final query row when qSeq < kvSeq and returns 3.75 instead of 1.5 there."
       },
       "attrs": { "num_heads": 1, "unidirectional": 1 },
       "inputs": {
@@ -662,7 +661,7 @@
     {
       "name": "dispatch_cliff_qseq_65537",
       "provenance": {
-        "notes": "query sequence 65537 exceeds the flash-path cap and exercises the scalar fallback's folded dispatch. x is clamped to 65535 and the grid-stride loop handles the final two queries on its second iteration. kvSeq=2 keeps the reference compact, while distinct query rows expose an indexing error."
       },
       "attrs": { "num_heads": 1 },
       "inputs": {
@@ -758,7 +757,7 @@
     {
       "name": "cluster_prefill_qseq64_h16_d64_compact",
       "provenance": {
-        "notes": "Compact companion for bench case prefill-cluster-qseq64-h16-d64-HEALTHY, which exposes the baseline/no-subgroups tier gap in the cluster prefill route."
       },
       "attrs": { "num_heads": 16, "scale": 0.125 },
       "inputs": {
@@ -785,7 +784,7 @@
     {
       "name": "flux2_joint_d128_q32_h8_compact",
       "provenance": {
-        "notes": "Compact companion for bench case flux2-joint-attn-img512-s1536-24h-d128, preserving headDim=128 and qSeq*heads>=256 so the flash-cluster path and no-subgroups fallback remain eligible."
       },
       "attrs": { "num_heads": 8, "scale": 0.08838834764831845 },
       "inputs": {
@@ -812,7 +811,7 @@
     {
       "name": "flux2_joint_d128_q32_h8_bias_compact",
       "provenance": {
-        "notes": "Compact correctness companion for the fused-bias benchmark at the 32-vec4, 32-lane register-topology boundary."
       },
       "attrs": { "num_heads": 8, "scale": 0.08838834764831845 },
       "inputs": {
@@ -869,7 +868,7 @@
     {
       "name": "q31_kv511_h8_d64_double_threshold_compact",
       "provenance": {
-        "notes": "Compact lock immediately below both flash admission boundaries: qSeq*heads=31*8=248 (<256) and kvSeq=511 (<512). It guards the realistic scalar fallback that the paired benchmark exposes."
       },
       "attrs": { "num_heads": 8, "scale": 0.125 },
       "inputs": {
@@ -886,7 +885,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 511, 512],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
       "outputs": {
@@ -913,7 +912,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.007, "cosStep": 0.041 }
         }
       },
       "outputs": {
@@ -938,7 +937,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": {
@@ -948,7 +947,7 @@
     {
       "name": "small_head_parallel_q3_kv65_h2_d6_compact",
       "provenance": {
-        "notes": "Compact correctness lock for the D<32 key-parallel fallback. The paired q128/kv2048 benchmark exposes the per-key-barrier pathology this path removes."
       },
       "attrs": { "num_heads": 2, "scale": 0.408248290463863 },
       "inputs": {
@@ -975,7 +974,7 @@
     {
       "name": "small_head_parallel_f16_q3_kv65_h2_d6",
       "provenance": {
-        "notes": "Float16 coverage for the D<32 key-parallel fallback. Values are scaled to O(1) while queries and keys remain small, keeping the convex-combination output near one so the tolerance detects meaningful scale errors."
       },
       "attrs": { "num_heads": 2, "scale": 0.408248290463863 },
       "inputs": {
@@ -1000,7 +999,7 @@
     {
       "name": "short_query_splitk_single_partition_q16_kv2048_h8_d32",
       "provenance": {
-        "notes": "Locks the adaptive one-partition short-query flash path once batch*queries*heads already supplies 128 workgroups; the paired q16/kv4096 benchmark measures the same occupancy regime."
       },
       "attrs": { "num_heads": 8, "scale": 0.17677669529663687 },
       "inputs": {
@@ -1017,7 +1016,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": {
@@ -1044,7 +1043,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": {
@@ -1079,7 +1078,7 @@
     {
       "name": "materialized_sgmat_fused_f16_q1024_h4_d80_no_bias",
       "provenance": {
-        "notes": "The d=80 twin of the f16 fused materialized fixture: a head dim that is not a multiple of 32 takes the guarded staging forms instead of the unguarded selects, on the same 16MB score-bytes admission floor."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
@@ -1104,7 +1103,7 @@
     {
       "name": "materialized_sgmat_q512_kv512_h1_d128_bias",
       "provenance": {
-        "notes": "Bias twin of the d128 subgroup-matrix prefill. headDim % 32 == 0 takes the select-form Q tile load, so this locks the fused Q bias on that branch and the post-normalization V bias in the compact store epilogue. The K bias is deliberately absent from the kernel and present in the reference: it shifts every logit in a row by one constant, so softmax cancels it, and a divergence here would prove otherwise."
       },
       "attrs": { "num_heads": 1, "scale": 0.08838834764831845 },
       "inputs": {
@@ -1136,7 +1135,7 @@
     {
       "name": "portable_head_dim320_no_bias",
       "provenance": {
-        "notes": "Locks the device-derived portable fallback above the former hard-coded headDim<=256 ceiling; D320 occurs in wide vision and multimodal attention blocks."
       },
       "attrs": { "num_heads": 1, "scale": 0.05590169943749474 },
       "inputs": {
@@ -1163,7 +1162,7 @@
     {
       "name": "splitk_decode_q1_kv1500_h6_d64_underfilled_wide_subgroup",
       "provenance": {
-        "notes": "A six-head f32 decode geometry where d64 supplies only 16 vec4 dot products to a 32-lane subgroup; the selector should prefer the portable split-K kernel instead of an underfilled cooperative QK subgroup."
       },
       "attrs": { "num_heads": 6, "scale": 0.125 },
       "inputs": {
@@ -1180,7 +1179,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1500, 384],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": {
@@ -1207,7 +1206,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1500, 128],
-          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.031, "cosStep": 0.019 }
         }
       },
       "outputs": {
@@ -1217,7 +1216,7 @@
     {
       "name": "materialized_prefill_q512_h1_d64_bias",
       "provenance": {
-        "notes": "Bias companion for the portable materialized f32 path; the even KV length exercises its register-cached vec4 softmax while locking fused Q bias and post-normalization V bias."
       },
       "attrs": { "num_heads": 1, "scale": 0.125 },
       "inputs": {
@@ -1281,7 +1280,7 @@
     {
       "name": "materialized_sgmat_cached_q520_kv516_h2_d80_all_tails",
       "provenance": {
-        "notes": "Generalized subgroup-matrix coverage: d80 exercises the BK32 tail, q520/kv516 exercise both matrix tile tails, and kv516 keeps the register-cached vec4 softmax eligible."
       },
       "attrs": { "num_heads": 2, "scale": 0.11180339887498948 },
       "inputs": {
@@ -1298,7 +1297,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 516, 160],
-          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.029, "cosStep": 0.017 }
         }
       },
       "outputs": {
@@ -1308,7 +1307,7 @@
     {
       "name": "materialized_sgmat_q520_kv516_h2_d80_all_tails_bias",
       "provenance": {
-        "notes": "Bias twin on the guarded Q tile load: d80 is not a multiple of 32, so the select-form branch is unavailable and the bounds-checked load carries the bias instead. q520/kv516 keep both matrix tile tails live, where a bias applied to a padded row would leak into the result."
       },
       "attrs": { "num_heads": 2, "scale": 0.11180339887498948 },
       "inputs": {
@@ -1357,7 +1356,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [1, 515, 80],
-          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.031, "cosStep": 0.019 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 513, 80], "tolerance": 0.001, "relTolerance": 0.001 } }
@@ -1365,7 +1364,7 @@
     {
       "name": "materialized_sgmat_cached_batch2_q512_kv512_h1_d32",
       "provenance": {
-        "notes": "Batch and narrow-head coverage for the generalized selector; two batches also exercise the workgroup-z to batch/head mapping."
       },
       "attrs": { "num_heads": 1, "scale": 0.17677669529663687 },
       "inputs": {
@@ -1382,7 +1381,7 @@
         "valueT": {
           "dtype": "float32",
           "shape": [2, 512, 32],
-          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.025, "cosStep": 0.013 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [2, 512, 32], "tolerance": 0.001, "relTolerance": 0.001 } }
@@ -1645,6 +1644,76 @@
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 4, 12], "tolerance": 0.00002 } }
     },
     {
       "name": "attn_bias_flash_f16_b1_h8_s32",
       "attrs": { "num_heads": 8 },
@@ -1662,7 +1731,7 @@
         "valueT": {
           "dtype": "float16",
           "shape": [1, 32, 256],
-          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
         },
         "attentionBiasT": {
           "dtype": "float16",
@@ -1675,7 +1744,7 @@
     {
       "name": "flash_q30_below_cluster_h9_d64_no_bias",
       "provenance": {
-        "notes": "One query row below the cooperative cluster threshold (qSeq 30 < 31) while qSeq*heads=270 still clears the flash occupancy floor, and kvSeq stays far under the split-K thresholds. Locks the one-query-row-per-workgroup flash kernel: the subgroup form on fixed-width tiers, the portable form on tiers without subgroups."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
@@ -1702,7 +1771,7 @@
     {
       "name": "flash_q30_below_cluster_h9_d64_bias",
       "provenance": {
-        "notes": "Fused-QKV-bias sibling of flash_q30_below_cluster_h9_d64_no_bias: same sub-cluster query count, so the biased flash kernel runs per query row instead of the clustered or split-K forms."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
@@ -1734,7 +1803,7 @@
     {
       "name": "flash_q30_below_cluster_h9_d64_attn_bias",
       "provenance": {
-        "notes": "Additive attention-bias sibling of flash_q30_below_cluster_h9_d64_no_bias: qSeq 30 keeps the tiled prefill kernel (qSeq >= 31) out, so the per-query-row attention-bias flash kernel runs on tiers without subgroups."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
@@ -1891,7 +1960,7 @@
     {
       "name": "materialized_sgmat_q512_kv512_h2_d64_offset_value_scale_lock",
       "provenance": {
-        "notes": "The other materialized-sgmat fixtures draw V from a zero-mean sinusoid, so a softmax average over hundreds of keys cancels to ~1e-5 and the 4e-4 absolute tolerance swamps any multiplicative error: halving the apply output passes every one of them. Offsetting V to oscillate about 1.0 makes the attention average O(1), so the normalization itself is under test - the denominator the row-stats pass publishes, and the divide the apply performs with it."
       },
       "attrs": { "num_heads": 2, "scale": 0.125 },
       "inputs": {
@@ -1918,7 +1987,7 @@
     {
       "name": "materialized_sgmat_fused_q1024_kv1024_h4_d64_offset_value_scale_lock",
       "provenance": {
-        "notes": "Above the fused-softmax score-byte threshold (4 heads x 1024 x 1024 x 4B = 16.8 MiB), so this is the row-stats route rather than the separate softmax pass. V oscillates about 1.0 so the attention average is O(1) and the normalization is actually under test: halving the apply output, or doubling either the per-tile or the combined denominator, fails this and passes every zero-mean fixture."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
@@ -1977,7 +2046,7 @@
     {
       "name": "splitk_decode_h4_d64_kv2048_no_bias_ramp_value_scale_lock",
       "provenance": {
-        "notes": "Every no-bias decode split-K fixture draws V from a zero-mean sinusoid, so the softmax average over 2048 keys cancels to ~5e-5 while the absolute tolerance is 3e-4: halving the apply output, or dropping the exp-weighted rescale in the cross-partition combine, passes all nine of them (the bias siblings only escape because the V bias survives the average). A monotone V ramp makes each output element 0.5 + 1.5 * (weighted mean key index) / kvSeq, which is O(1) - so a uniform scale error is visible - and is a direct function of the per-partition weights, so a combine that mixes partitions with the wrong exponent offset moves it too."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
@@ -2002,7 +2071,7 @@
     {
       "name": "splitk_decode_h4_d64_kv2048_no_bias_f16_ramp_value_scale_lock",
       "provenance": {
-        "notes": "The f16 arm of the no-bias decode split-K kernel is a separate compile (usesF16) and its only coverage was flash_decode_head_dim64_no_bias_f16, whose 0.03 absolute tolerance against a ~1e-4 expected output cannot detect a 295x scale error - the single blindest fixture in the op. Same ramped V as the f32 sibling, so the f16 accumulate/rescale path is scale-locked too."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
@@ -2027,7 +2096,7 @@
     {
       "name": "flash_q32_h4_d128_no_bias_f16_ramp_value_scale_lock",
       "provenance": {
-        "notes": "qkv_no_bias_flash_q32_broadcast and qkv_no_bias_flash_q32_shared are reached by exactly one fixture, flash_self_attention_head_dim128_no_bias_f16, whose zero-mean V averages to ~1.8e-3 under a 0.03 absolute tolerance - a 16x scale error passes. Same shape, dtype and route with a monotone V ramp so each query row lands at 0.5 + 1.5 * (weighted mean key index) / 64: the register-blocked q32 epilogue divide is now under test, and the per-query variation also catches a broadcast that staged the wrong lane's row."
       },
       "attrs": { "num_heads": 4, "scale": 0.08838834764831845 },
       "inputs": {
@@ -2052,7 +2121,7 @@
     {
       "name": "short_query_splitk_single_partition_q16_kv2048_h8_d32_ramp_value_scale_lock",
       "provenance": {
-        "notes": "The adaptive short-query path uses one partition when batch*queries*heads already supplies enough work, so normalization stays in-kernel instead of using the cross-partition combine. A ramped V keeps each of the 16 query-row outputs O(1), making scale errors observable where a zero-mean companion averages to 5.8e-5 under a 5e-4 tolerance."
       },
       "attrs": { "num_heads": 8, "scale": 0.17677669529663687 },
       "inputs": {
@@ -2073,6 +2142,179 @@
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 16, 256], "tolerance": 0.00005 } }
     }
   ]
 }

 {
   "cases": [
     {
       "name": "cross_attention_no_bias",
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/multihead_attention_op_test.cc",
         "test": "MultiHeadAttentionTest.CrossAttention_Batch1_HeadSize16",
+        "notes": "ORT cross-attention with batch 1, two query rows, three key/value rows, two heads of width 16, and no bias."
       },
       "attrs": { "num_heads": 2 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 64, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.0007669903939 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 64, 256], "tolerance": 0.0002 } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00009587379924 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0002 } }
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00002396844981 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0003 } }
     {
       "name": "splitk_decode_head_dim128_no_bias_kv2048",
       "provenance": {
+        "notes": "A compact decode with eight heads, head size 128, and a long KV sequence exercises split-K attention."
       },
       "attrs": { "num_heads": 2, "scale": 0.08838834764831845 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00002396844981 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 1, 256], "tolerance": 0.0003 } }
         "valueT": {
           "dtype": "float32",
           "shape": [2, 768, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.0122718463, "cosStep": 0.00001597896654 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [2, 1, 512], "tolerance": 0.0003 } }
     {
       "name": "flash_self_attention_head_dim128_no_bias_f16",
       "provenance": {
+        "notes": "A float16 prefill with head size 128 exercises the q32 broadcast route with subgroups and the shared-memory q32 route without them."
       },
       "attrs": { "num_heads": 4, "scale": 0.08838834764831845 },
       "inputs": {
         "valueT": {
           "dtype": "float16",
           "shape": [1, 64, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.0122718463, "cosStep": 0.000383495197 }
         }
       },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 64, 512], "tolerance": 0.002 } }
     },
     {
       "name": "flash_decode_head_dim64_no_bias_f16",
         "valueT": {
           "dtype": "float16",
           "shape": [1, 512, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00009587379924 }
         }
       },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1, 256], "tolerance": 0.002 } }
     },
     {
       "name": "cross_attention_projection_bias_f16",
       "name": "causal_cross_attention_q2_kv4_exact",
       "provenance": {
         "source": "onnxruntime/contrib_ops/cpu/bert/attention_helper.h",
+        "notes": "Upper-left causal masking with two queries and four keys. The second query averages only the first two values, producing exactly 1.5."
       },
       "attrs": { "num_heads": 1, "unidirectional": 1 },
       "inputs": {
     {
       "name": "dispatch_cliff_qseq_65537",
       "provenance": {
+        "notes": "A 65,537-row query exceeds the flash-path cap and exercises the scalar fallback's folded dispatch. The x dimension is capped at 65,535, so a grid-stride iteration handles the final two distinct rows; two keys keep expected-value evaluation compact."
       },
       "attrs": { "num_heads": 1 },
       "inputs": {
     {
       "name": "cluster_prefill_qseq64_h16_d64_compact",
       "provenance": {
+        "notes": "A compact sequence-64, 16-head prefill exercises the clustered route and its no-subgroups fallback."
       },
       "attrs": { "num_heads": 16, "scale": 0.125 },
       "inputs": {
     {
       "name": "flux2_joint_d128_q32_h8_compact",
       "provenance": {
+        "notes": "Head size 128 and `qSeq * heads >= 256` keep the flash-cluster route and its no-subgroups fallback eligible at a compact scale."
       },
       "attrs": { "num_heads": 8, "scale": 0.08838834764831845 },
       "inputs": {
     {
       "name": "flux2_joint_d128_q32_h8_bias_compact",
       "provenance": {
+        "notes": "A compact fused-bias case exercises the boundary of 32 four-wide head segments across 32 lanes."
       },
       "attrs": { "num_heads": 8, "scale": 0.08838834764831845 },
       "inputs": {
     {
       "name": "q31_kv511_h8_d64_double_threshold_compact",
       "provenance": {
+        "notes": "`qSeq * heads = 248` and `kvSeq = 511` lie immediately below the two flash admission boundaries, selecting the scalar fallback."
       },
       "attrs": { "num_heads": 8, "scale": 0.125 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 511, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.0122718463, "cosStep": 0.0000480307096 }
         }
       },
       "outputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.04908738521, "cosStep": 0.00004793689962 }
         }
       },
       "outputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.04908738521, "cosStep": 0.0001917475985 }
         }
       },
       "outputs": {
     {
       "name": "small_head_parallel_q3_kv65_h2_d6_compact",
       "provenance": {
+        "notes": "Head size 6 exercises the key-parallel fallback with multiple keys per lane and a compact expected output."
       },
       "attrs": { "num_heads": 2, "scale": 0.408248290463863 },
       "inputs": {
     {
       "name": "small_head_parallel_f16_q3_kv65_h2_d6",
       "provenance": {
+        "notes": "Float16 head size 6 exercises the key-parallel fallback. Values remain near one while queries and keys stay small, keeping scale errors observable."
       },
       "attrs": { "num_heads": 2, "scale": 0.408248290463863 },
       "inputs": {
     {
       "name": "short_query_splitk_single_partition_q16_kv2048_h8_d32",
       "provenance": {
+        "notes": "Sixteen queries over eight heads supply 128 workgroups, so the adaptive short-query flash path uses one K partition and normalizes in-kernel."
       },
       "attrs": { "num_heads": 8, "scale": 0.17677669529663687 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 2048, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.02454369261, "cosStep": 0.00002396844981 }
         }
       },
       "outputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.04908738521, "cosStep": 0.0001917475985 }
         }
       },
       "outputs": {
     {
       "name": "materialized_sgmat_fused_f16_q1024_h4_d80_no_bias",
       "provenance": {
+        "notes": "Float16 head size 80 is not divisible by 32, selecting guarded tile staging on the fused materialized route at the 16 MiB score threshold."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
     {
       "name": "materialized_sgmat_q512_kv512_h1_d128_bias",
       "provenance": {
+        "notes": "Head size 128 selects the subgroup-matrix prefill's unguarded Q tile load. Q bias is fused into that load and V bias follows normalization; K bias adds one constant to every row logit and must cancel under softmax."
       },
       "attrs": { "num_heads": 1, "scale": 0.08838834764831845 },
       "inputs": {
     {
       "name": "portable_head_dim320_no_bias",
       "provenance": {
+        "notes": "Head size 320 exercises the portable route used by wide vision and multimodal attention blocks."
       },
       "attrs": { "num_heads": 1, "scale": 0.05590169943749474 },
       "inputs": {
     {
       "name": "splitk_decode_q1_kv1500_h6_d64_underfilled_wide_subgroup",
       "provenance": {
+        "notes": "A six-head f32 decode geometry where head size 64 supplies only 16 vec4 dot products to a 32-lane subgroup. The case exercises the portable split-K route instead of an underfilled cooperative QK subgroup."
       },
       "attrs": { "num_heads": 6, "scale": 0.125 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1500, 384],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.01636246174, "cosStep": 0.00002181661565 }
         }
       },
       "outputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 1500, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.04908738521, "cosStep": 0.00006544984695 }
         }
       },
       "outputs": {
     {
       "name": "materialized_prefill_q512_h1_d64_bias",
       "provenance": {
+        "notes": "An even key/value length exercises register-cached vec4 softmax on the portable materialized route with fused Q bias and post-normalization V bias."
       },
       "attrs": { "num_heads": 1, "scale": 0.125 },
       "inputs": {
     {
       "name": "materialized_sgmat_cached_q520_kv516_h2_d80_all_tails",
       "provenance": {
+        "notes": "Head size 80 exercises the BK32 tail, query length 520 and key/value length 516 exercise both matrix tile tails, and the even key count keeps register-cached vec4 softmax eligible."
       },
       "attrs": { "num_heads": 2, "scale": 0.11180339887498948 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 516, 160],
+          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.03926990817, "cosStep": 0.0001522089464 }
         }
       },
       "outputs": {
     {
       "name": "materialized_sgmat_q520_kv516_h2_d80_all_tails_bias",
       "provenance": {
+        "notes": "Head size 80 selects the bounds-checked Q tile load with fused bias. Query length 520 and key/value length 516 keep both matrix tails live, exposing any bias applied to padded rows."
       },
       "attrs": { "num_heads": 2, "scale": 0.11180339887498948 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [1, 515, 80],
+          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.07853981634, "cosStep": 0.0003050089955 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 513, 80], "tolerance": 0.001, "relTolerance": 0.001 } }
     {
       "name": "materialized_sgmat_cached_batch2_q512_kv512_h1_d32",
       "provenance": {
+        "notes": "Two batches and head size 32 exercise the narrow-head subgroup-matrix route and the workgroup-z mapping to batch and head."
       },
       "attrs": { "num_heads": 1, "scale": 0.17677669529663687 },
       "inputs": {
         "valueT": {
           "dtype": "float32",
           "shape": [2, 512, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.06, "sinStep": 0.1963495408, "cosStep": 0.000383495197 }
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [2, 512, 32], "tolerance": 0.001, "relTolerance": 0.001 } }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 4, 12], "tolerance": 0.00002 } }
     },
+    {
+      "name": "attn_bias_scalar_all_neg_inf_row_zero_q2_kv32_h8_d64",
+      "provenance": {
+        "notes": "BSH float32 Q/K/V uses eight heads of width 64, two query tokens, and 32 keys. Rank-4 attention bias is broadcast over batch and heads; query row 1 is -Infinity at every key while row 0 is finite. The portable scalar online-attention route must return zero for every head in row 1 without forming NaN from a zero denominator; row 0 follows the ordinary softmax-weighted value calculation."
+      },
+      "attrs": { "num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 2, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.011, "cosStep": 0.023 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.007, "cosStep": 0.041 }
+        },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 1, 2, 32],
+          "data": {
+            "kind": "values",
+            "values": [0.0, 0.193, 0.296, 0.259, 0.1, -0.105, -0.261, -0.295, -0.189, 0.005, 0.197, 0.296, 0.256, 0.096, -0.11, -0.264, -0.294, -0.185, 0.01, 0.201, 0.297, 0.254, 0.091, -0.115, -0.266, -0.293, -0.181, 0.015, 0.205, 0.298, 0.251, 0.086, "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity"]
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 2, 512], "tolerance": 0.00002, "relTolerance": 0.00002 }
+      }
+    },
+    {
+      "name": "attn_bias_scalar_single_finite_key_q1_kv32_h8_d64",
+      "provenance": {
+        "notes": "Rank-4 bias has one finite entry at key 3 and -Infinity elsewhere. The scalar softmax therefore places all mass on key 3, so every head returns value row 3. This distinguishes a valid one-key denominator from `attn_bias_scalar_all_neg_inf_row_zero_q2_kv32_h8_d64`."
+      },
+      "attrs": { "num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.011, "cosStep": 0.023 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.007, "cosStep": 0.041 }
+        },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1, 32],
+          "data": {
+            "kind": "values",
+            "values": ["-Infinity", "-Infinity", "-Infinity", 0.0, "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity", "-Infinity"]
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 512], "tolerance": 0.00001, "relTolerance": 0.00001 }
+      }
+    },
     {
       "name": "attn_bias_flash_f16_b1_h8_s32",
       "attrs": { "num_heads": 8 },
         "valueT": {
           "dtype": "float16",
           "shape": [1, 32, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.02454369261, "cosStep": 0.001533980788 }
         },
         "attentionBiasT": {
           "dtype": "float16",
     {
       "name": "flash_q30_below_cluster_h9_d64_no_bias",
       "provenance": {
+        "notes": "Query length 30 is below the cooperative-cluster threshold while 270 query-head rows clear the flash admission floor. The case exercises one-query-row-per-workgroup flash with and without subgroups."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
     {
       "name": "flash_q30_below_cluster_h9_d64_bias",
       "provenance": {
+        "notes": "Fused QKV bias with query length 30 exercises the biased one-query-row flash route below the clustered and split-K thresholds."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
     {
       "name": "flash_q30_below_cluster_h9_d64_attn_bias",
       "provenance": {
+        "notes": "Additive attention bias with query length 30 exercises the per-query-row flash route below the tiled-prefill threshold on tiers without subgroups."
       },
       "attrs": { "num_heads": 9, "scale": 0.125 },
       "inputs": {
     {
       "name": "materialized_sgmat_q512_kv512_h2_d64_offset_value_scale_lock",
       "provenance": {
+        "notes": "Values oscillate around 1.0, keeping the attention average at order-one magnitude. This makes both the denominator published by row statistics and the division in the subgroup-matrix apply pass observable."
       },
       "attrs": { "num_heads": 2, "scale": 0.125 },
       "inputs": {
     {
       "name": "materialized_sgmat_fused_q1024_kv1024_h4_d64_offset_value_scale_lock",
       "provenance": {
+        "notes": "The 16.8 MiB score tensor selects the fused row-statistics route. Values oscillate around 1.0 so a per-tile, combined-denominator, or apply-scale error remains observable."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
     {
       "name": "splitk_decode_h4_d64_kv2048_no_bias_ramp_value_scale_lock",
       "provenance": {
+        "notes": "A monotone value ramp makes each output an order-one function of the weighted key index. This exposes both uniform scale errors and incorrect exponential offsets in the cross-partition split-K combine."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
     {
       "name": "splitk_decode_h4_d64_kv2048_no_bias_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 final normalization observable."
       },
       "attrs": { "num_heads": 4, "scale": 0.125 },
       "inputs": {
     {
       "name": "flash_q32_h4_d128_no_bias_f16_ramp_value_scale_lock",
       "provenance": {
+        "notes": "A monotone value ramp gives each float16 q32 row a distinct order-one output, exposing the register-blocked epilogue division and an incorrect broadcast lane."
       },
       "attrs": { "num_heads": 4, "scale": 0.08838834764831845 },
       "inputs": {
     {
       "name": "short_query_splitk_single_partition_q16_kv2048_h8_d32_ramp_value_scale_lock",
       "provenance": {
+        "notes": "The adaptive short-query path uses one partition because query-head parallelism already supplies enough work. A value ramp keeps all 16 outputs at order-one magnitude, exposing in-kernel normalization errors."
       },
       "attrs": { "num_heads": 8, "scale": 0.17677669529663687 },
       "inputs": {
         }
       },
       "outputs": { "outputT": { "dtype": "float32", "shape": [1, 16, 256], "tolerance": 0.00005 } }
+    },
+    {
+      "name": "small_seq_blocked_bias_causal_1x64x64_h4d16",
+      "provenance": {
+        "notes": "Short-sequence attention with 64 queries, 64 keys, four heads of width 16, fused bias, and causal masking."
+      },
+      "attrs": { "num_heads": 4, "unidirectional": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [192],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "small_seq_blocked_no_bias_causal_q13_kv40_h2d16",
+      "provenance": {
+        "notes": "Thirteen queries leave five active rows in the final eight-row block, and the key count is not divisible by sixteen. Causal masking makes out-of-range lanes and inactive rows contribute the identity partial."
+      },
+      "attrs": { "num_heads": 2, "unidirectional": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 13, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 40, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 40, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 13, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "small_seq_blocked_bias_bidirectional_q24_kv40_h3d32",
+      "provenance": {
+        "notes": "Head dim 32 (eight vec4 per row, the route's register ceiling) with bias and no causal mask; kv > q so every query sees all forty keys."
+      },
+      "attrs": { "num_heads": 3, "unidirectional": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 24, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 40, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 40, 96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [288],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 24, 96], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "small_seq_blocked_no_bias_bidirectional_q3_kv5_h1d8",
+      "provenance": {
+        "notes": "Fewer keys than key lanes (5 keys, 8 lanes after the power-of-two floor) and fewer queries than one block: most lanes hold no key and merge identities."
+      },
+      "attrs": { "num_heads": 1, "unidirectional": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 5, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 5, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [1, 3, 8], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "small_seq_blocked_bias_causal_batch2_q16_kv16_h2d16",
+      "provenance": {
+        "notes": "Two batches of two heads at exactly two query blocks; the batch axis rides dispatch z and the bias is shared across batches."
+      },
+      "attrs": { "num_heads": 2, "unidirectional": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 16, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [2, 16, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 16, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.13 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float32", "shape": [2, 16, 32], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "small_seq_positive_infinity_query",
+      "provenance": {
+        "notes": "The unmasked whole-head route always includes an `exp(0)` term, so a finite row cannot have a zero denominator. A positive-infinity query produces non-finite scores and the expected NaN result."
+      },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8],
+          "data": { "kind": "cycle", "values": ["Infinity", 0.5, -0.25, 0.75, 0.125, -0.5, 0.25, 1.0] }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.023, "scale": 0.4 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.007, "cosStep": 0.041, "scale": 0.4 }
+        }
+      },
+      "outputs": {
+        "outputT": {
+          "dtype": "float32",
+          "shape": [1, 4, 8],
+          "tolerance": 0.00002,
+          "relTolerance": 0.00002,
+          "allowNaN": true
+        }
+      }
     }
   ]
 }