sync 2e7068faf55e

Browse files

Files changed (15) hide show

README.md +129 -0
build/webgpu/attention-rank4-tiled.wgsl.jinja +119 -0
build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja +131 -0
build/webgpu/attn-flash-decode-splitk.wgsl.jinja +549 -0
build/webgpu/attn-flash-online.wgsl.jinja +367 -0
build/webgpu/attn-flash-prefill-cluster.wgsl.jinja +554 -0
build/webgpu/attn-flash-q32-broadcast.wgsl.jinja +246 -0
build/webgpu/attn-online-scalar.wgsl.jinja +205 -0
build/webgpu/bench.json +1565 -0
build/webgpu/gqa-attention.wgsl.jinja +285 -0
build/webgpu/gqa-present.wgsl.jinja +440 -0
build/webgpu/gqa-qprep.wgsl.jinja +57 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +27 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.GroupQueryAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Grouped-query attention for explicit BSH Q/K/V and BNSH caches. Direct Q/K/V supports bidirectional attention or causal local windows and may store its generated float cache independently as float16 or float32; existing unquantized cache inputs match the Q/K/V dtype. Causal cache paths support rotary embeddings, sliding windows, bias, head sinks, softcap, smooth softmax, and paired Q/K RMS normalization. Int8/int4 caches require float32 Q/K/V and output; int4 is prompt-only. Packed QKV, position IDs, interleaved rotary, bfloat16/float8, and diagnostic QK output are not implemented.
+See the [ONNX Runtime `GroupQueryAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.GroupQueryAttention) 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, num_heads * head_size)`. | required |
+| `key` | `keyT` | `T` | `3` | — | Key tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
+| `value` | `valueT` | `T` | `3` | — | Value tensor of shape `(batch_size, kv_sequence_length, kv_num_heads * head_size)`. | required |
+| `past_key` | `pastKeyT` | `T_CACHE` | `4` | — | Optional cached key state in BNSH format. Its sequence axis is `max_sequence_length` when the past and present buffers are shared, otherwise `past_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | optional |
+| `past_value` | `pastValueT` | `T_CACHE` | `4` | — | Optional cached value state in BNSH format with the same length and packing semantics as `past_key`. | optional |
+| `seqlens_k` | `seqlensKT` | `M` | `1` | — | Int32 tensor of shape `(batch_size)` containing each sample's total sequence length minus one. | required |
+| `total_sequence_length` | `totalSequenceLengthT` | `M` | `1` | — | Length-one int32 tensor containing the maximum total sequence length (past plus new) in the batch. | required |
+| `cos_cache` | `cosCacheT` | `T` | `2` | — | Optional cosine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
+| `sin_cache` | `sinCacheT` | `T` | `2` | — | Optional sine cache for rotary embeddings with shape `(max_sequence_length, head_size / 2)`. | optional |
+| `attention_bias` | `attentionBiasT` | `T` | `4` | — | Optional additive term for QK scores with shape `(batch_size or 1, num_heads or 1, sequence_length, total_sequence_length)`; the first two dimensions broadcast. | optional |
+| `head_sink` | `headSinkT` | `T` | `1` | — | Optional per-head smooth factor of shape `(num_heads)` added to the softmax denominator. | optional |
+| `k_scale` | `kScaleT` | `T_KV_SCALE` | `1` | — | Optional float32 key-cache scale: one value for `PER_TENSOR`, or `kv_num_heads * head_size` values for `PER_CHANNEL`. | optional |
+| `v_scale` | `vScaleT` | `T_KV_SCALE` | `1` | — | Optional float32 value-cache scale with the same shape convention as `k_scale`. | optional |
+| `q_norm_weight` | `qNormWeightT` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to queries before rotary embedding. It must be provided together with `k_norm_weight`. | optional |
+| `k_norm_weight` | `kNormWeightT` | `T` | `1` | — | Optional per-head RMS-normalization weight of shape `(head_size)` applied to keys before rotary embedding. It must be provided together with `q_norm_weight`. | optional |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `outputT` | `T` | `3` | same as `query` | Attention output of shape `(batch_size, sequence_length, hidden_size)`. | required |
+| `present_key` | `presentKeyT` | `T_CACHE` | `4` | — | Updated key cache in BNSH format. Its sequence axis is `max_sequence_length` for a shared buffer, otherwise `past_sequence_length + kv_sequence_length`; int4 stores each signed value as a +8-biased nibble, with the even head coordinate low, packing two values per logical uint8 element and widening each byte to one u32 WebGPU buffer word. | required |
+| `present_value` | `presentValueT` | `T_CACHE` | `4` | — | Updated value cache in BNSH format with the same length and packing semantics as `present_key`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `causal` | `1` | Whether to apply a causal mask. Set to 0 for bidirectional attention; `local_window_size` must then be -1. |
+| `do_rotary` | `0` | Set to 1 to apply rotary position embeddings. The default 0 disables them. |
+| `k_quant_type` | `"NONE"` | Key-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
+| `local_window_size` | `-1` | Left window size for causal local attention. The default -1 disables local attention, and the value must be -1 when `causal` is 0. |
+| `qk_norm_epsilon` | `0.000001` | Epsilon for the per-head Q/K RMS normalization applied when both normalization weights are provided. |
+| `sliding_window_cache` | `0` | Set to 1 when past/present caches are fixed-size window buffers that evict old tokens from the front. Requires `local_window_size > 0` and enough cache capacity. |
+| `smooth_softmax` | `-1` | Set to 1 to enable the smooth-softmax denominator term. |
+| `softcap` | `0` | Positive softcap applied to attention scores. The default 0 disables soft-capping. |
+| `v_quant_type` | `"NONE"` | Value-cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
+| `kv_cache_bit_width` | — | Quantized cache bit width, either 8 or 4. Four-bit values are packed two per uint8 element. |
+| `kv_num_heads` | — | Number of key/value attention heads. |
+| `num_heads` | — | Number of query attention heads. |
+| `scale` | — | Optional QK score scale; zero or omission selects `1 / sqrt(head_size)`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `T_CACHE` | `float32`, `float16`, `uint8`, `int8` |
+| `T_KV_SCALE` | `float32` |
+| `M` | `int32` |
+## Device requirements
+Some implementation variants require `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
+- [`attention-rank4-tiled.wgsl.jinja`](build/webgpu/attention-rank4-tiled.wgsl.jinja)
+- [`attn-flash-decode-splitk-merge.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk-merge.wgsl.jinja)
+- [`attn-flash-decode-splitk.wgsl.jinja`](build/webgpu/attn-flash-decode-splitk.wgsl.jinja)
+- [`attn-flash-online.wgsl.jinja`](build/webgpu/attn-flash-online.wgsl.jinja)
+- [`attn-flash-prefill-cluster.wgsl.jinja`](build/webgpu/attn-flash-prefill-cluster.wgsl.jinja)
+- [`attn-flash-q32-broadcast.wgsl.jinja`](build/webgpu/attn-flash-q32-broadcast.wgsl.jinja)
+- [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
+- [`gqa-attention.wgsl.jinja`](build/webgpu/gqa-attention.wgsl.jinja)
+- [`gqa-present.wgsl.jinja`](build/webgpu/gqa-present.wgsl.jinja)
+- [`gqa-qprep.wgsl.jinja`](build/webgpu/gqa-qprep.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
+The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
+- `presentKeyT`
+- `presentValueT`
+Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.GroupQueryAttention", { version: 1 });
+// Explicit destinations request optional results or supply metadata that cannot be inferred.
+const { outputT, presentKeyT, presentValueT } = await kernel({
+  queryT: { data: queryTData, shape: [2, 1, 16] },
+  keyT: { data: keyTData, shape: [2, 1, 8] },
+  valueT: { data: valueTData, shape: [2, 1, 8] },
+  pastKeyT: { data: pastKeyTData, shape: [2, 1, 8, 8] },
+  pastValueT: { data: pastValueTData, shape: [2, 1, 8, 8] },
+  seqlensKT: { data: seqlensKTData, shape: [2] },
+  totalSequenceLengthT: { data: totalSequenceLengthTData, shape: [1] },
+}, {
+  attrs: { num_heads: 2, kv_num_heads: 1 },
+  outputs: {
+    presentKeyT: { shape: [2, 1, 8, 8], dtype: "float32" },
+    presentValueT: { shape: [2, 1, 8, 8], dtype: "float32" },
+  },
+});
+```

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

	@@ -0,0 +1,119 @@

+{% 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;
+var<workgroup> acc: array<f32, {{ vHeadCap * blockM }}>;
+{% set ATTN_SCALE_DIM = "params.headSize" %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+fn kv_head(q_head: u32) -> u32 {
+  return q_head / (params.qHeads / params.kvHeads);
+}
+@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.
+  let batchHeads = params.count / (params.vHeadSize * params.qSeq);
+  if (tile >= batchHeads * qTiles) { return; }
+  let qTile = tile % qTiles;
+  var t = tile / qTiles;
+  let qh = t % params.qHeads;
+  let batch = t / params.qHeads;
+  let qs = qTile * BLOCK_M + tid;
+  // Threads past the q-seq tail own no query row. No barriers below, so they just
+  // exit (they hold no shared state the surviving threads read).
+  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) {
+    acc[d * BLOCK_M + tid] = 0.0;
+  }
+  var running_max: f32 = -3.4028234663852886e38;
+  var running_denom: f32 = 0.0;
+  // Causal: query qs attends keys [0, qs]. Per-thread bound — threads in the tile
+  // diverge on trip count, which is fine (no barriers).
+  var maxKj = params.kvSeq;
+  if (params.isCausal != 0u) {
+    maxKj = min(maxKj, qs + 1u);
+  }
+  for (var ks: u32 = 0u; ks < maxKj; ks = ks + 1u) {
+    var masked = false;
+{% if hasMask and maskIsBool %}
+    let mIdxB = batch * params.maskBatchStride + qh * params.maskHeadStride + qs * params.maskSeqStride + ks;
+    masked = attn_mask[mIdxB] == 0u;
+{% endif %}
+    // A rejected bool-mask key has zero softmax mass. Skipping it is safe here:
+    // this thread-per-query kernel has no barriers inside the key loop.
+    if (masked) { continue; }
+    // 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]);
+      }
+      score = dot * scale;
+      if (params.softcap != 0.0) {
+        score = params.softcap * tanh(score / params.softcap);
+      }
+{% if hasMask and not maskIsBool %}
+      let mIdxA = batch * params.maskBatchStride + qh * params.maskHeadStride + qs * params.maskSeqStride + ks;
+      score = score + f32(attn_mask[mIdxA]);
+{% endif %}
+    }
+    let next_max = max(running_max, score);
+    let prev_scale = exp(running_max - next_max);
+    let weight = exp(score - next_max);
+    running_max = next_max;
+    running_denom = running_denom * prev_scale + weight;
+    let vRow = (batch * params.kvSeq + ks) * params.vHidden + kh * params.vHeadSize;
+    for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
+      acc[d * BLOCK_M + tid] = acc[d * BLOCK_M + tid] * prev_scale + weight * f32(v[vRow + d]);
+    }
+  }
+  let inv_denom = select(0.0, 1.0 / running_denom, running_denom > 0.0);
+  let yBase = (batch * params.qSeq + qs) * (params.qHeads * params.vHeadSize) + qh * params.vHeadSize;
+  for (var d: u32 = 0u; d < params.vHeadSize; d = d + 1u) {
+    y[yBase + d] = {{ scalar }}(acc[d * BLOCK_M + tid] * inv_denom);
+  }
+}

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

	@@ -0,0 +1,131 @@

+{% if scalar == "f16" %}
+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
+// states (max, denom, weighted V sum) written by the split pass into the final
+// attention output for each query token, using the online softmax merge:
+//   globalMax  = max_s partialMax[s]
+//   globalDen  = sum_s partialDenom[s] * exp(partialMax[s] - globalMax)
+//   out[d]     = (sum_s partialOut[s, d] * exp(partialMax[s] - globalMax)) / globalDen
+// One workgroup per (batch, head); one thread per output vec4 channel. The
+// per-split (max, denom) scan is tiny (NUM_SPLITS is small), so every thread
+// recomputes it locally and keeps the merge state in registers.
+{% if source.layout == "bhsd" %}
+// Output layout: rank-4 [batch, heads, seq, headDim], matching the split pass.
+{% elif source.layout == "layer_cache" %}
+// Output layout: flat [heads, headDim], matching Qwen3.5 decode.
+{% else %}
+// Output layout: token-major [batch, seq, heads * headDim], matching the split pass.
+{% endif %}
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const NUM_SPLITS: u32 = {{ numSplits }}u;
+{% if splitQueries %}
+const Q_SEQ: u32 = {{ qSeq }}u;
+{% endif %}
+{% if source.layout == "bsh" %}
+const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
+{% endif %}
+{% if hasBias %}
+// Packed [Q; K; V] bias rows. The V bias is token-independent, so the split
+// pass leaves it off and it is added here once after the final normalize,
+// matching the non-split flash-decode semantics exactly.
+const HEAD_DIM: u32 = {{ headDim }}u;
+const Q_HIDDEN: u32 = {{ qHidden }}u;
+{% endif %}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+@compute @workgroup_size(HEAD_DIM_V4, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+{% if splitQueries %}
+  let queryToken = wg.x;
+{% endif %}
+  let h = wg.y;
+  let b = wg.z;
+  let d4 = lid.x;
+  if (h >= Q_HEADS || d4 >= HEAD_DIM_V4{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) {
+    return;
+  }
+{% if splitQueries %}
+  let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS;
+{% else %}
+  let mdBase = (b * Q_HEADS + h) * NUM_SPLITS;
+{% endif %}
+  var globalMax = -FLT_MAX;
+  for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
+    globalMax = max(globalMax, partial_stats[mdBase + s].x);
+  }
+  var globalDenom = 0.0;
+  var outV = vec4<f32>(0.0);
+  for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) {
+    let stats = partial_stats[mdBase + s];
+    let w = exp_shift(stats.x, globalMax);
+    globalDenom = globalDenom + stats.y * w;
+{% if splitQueries %}
+    let pBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
+{% else %}
+    let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4;
+{% endif %}
+    outV = outV + partial_out[pBase + d4] * w;
+  }
+{% if source.layout == "bsh" %}
+{% if splitQueries %}
+  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% endif %}
+{% elif source.layout == "layer_cache" %}
+  let qBaseV4 = h * HEAD_DIM_V4;
+{% else %}
+{% if splitQueries %}
+  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
+{% endif %}
+{% endif %}
+  // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty —
+  // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1,
+  // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty
+  // row has a defined zero result.
+  var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0);
+{% if hasBias %}
+  // V bias row base: 2 * qHidden (skip the packed Q and K bias blocks) + this head.
+  let vBiasBase = 2u * Q_HIDDEN + h * HEAD_DIM + d4 * 4u;
+  outValue = outValue + vec4<f32>(bias[vBiasBase], bias[vBiasBase + 1u], bias[vBiasBase + 2u], bias[vBiasBase + 3u]);
+{% endif %}
+{% if hasGate %}
+  // 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 %}
+  output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
+}

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

	@@ -0,0 +1,549 @@

+{% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
+{% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
+{% if quantizedCache is not defined %}{% set quantizedCache = false %}{% endif %}
+{% if cacheSeqlens is not defined %}{% set cacheSeqlens = false %}{% endif %}
+{% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
+{% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
+{% set 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
+// KV sequence into `NUM_SPLITS` contiguous ranges and gives each range its own
+// workgroup, so `batch * numHeads * NUM_SPLITS` workgroups run the tiled online
+// softmax in parallel. Each workgroup emits the *un-normalized* online state for
+// its range — the running (max, denom) and the softmax-weighted V sum before the
+// final divide — and the merge pass combines the per-split states with the online
+// rule.
+{% if source.layout == "bhsd" %}
+// Layout: rank-4 [batch, heads, seq, headDim] for Q/K/V.
+{% elif source.layout == "layer_cache" %}
+// Layout: flat query [heads, headDim] plus a persistent KV cache laid out
+// [layer, cacheLen, kvHeads, headDim]. This is the Qwen3.5 decode layout; the
+// dispatch has a single implicit batch.
+{% else %}
+// Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
+// are compiled constants.
+{% endif %}
+{% if (fusedQNormRope is defined and fusedQNormRope) or source.layout != "layer_cache" %}const HEAD_DIM: u32 = {{ headDim }}u;
+{% endif %}
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+{% if source.layout == "bsh" %}
+const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
+const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
+{% elif source.layout == "layer_cache" %}
+const LAYER: u32 = {{ layer }}u;
+const CACHE_LEN: u32 = {{ cacheLen }}u;
+const ATTN_SCALE: f32 = {{ scale }};
+{% endif %}
+const WG: u32 = {{ splitKWorkgroupSize }}u;
+const NUM_SPLITS: u32 = {{ numSplits }}u;
+{% if splitQueries %}
+const Q_SEQ: u32 = {{ qSeq }}u;
+{% endif %}
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+var<workgroup> q_shared: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> running_out: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> probs: array<f32, WG>;
+{% set coopQk = useSubgroups and headDimV4 >= 8 and not (usesF16 and headDimV4 <= 32) %}
+{% set jGroups = (splitKWorkgroupSize / headDimV4)|int %}
+{% set jSplitV = (splitKWorkgroupSize % headDimV4 == 0) and (jGroups >= 2) %}
+{% if coopQk %}
+var<workgroup> sval_sh: array<f32, WG>;
+{% endif %}
+{% if jSplitV %}
+var<workgroup> vacc_sh: array<vec4<f32>, WG>;
+{% endif %}
+{% set combineSubgroups = useSubgroups %}
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// 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
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  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);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. 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;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+{% if source.layout == "layer_cache" %}{% set ATTN_SCALE_OVERRIDE = "ATTN_SCALE" %}{% endif %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+{% if ATTN_SCALE_OVERRIDE is defined %}
+  return {{ ATTN_SCALE_OVERRIDE }};
+{% else %}
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+{% endif %}
+}
+{% if quantizedCache %}
+{% macro emit_quant_scale4(kind, scaleBuffer) %}
+fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
+  if (params.perChannel == 0u) {
+    return vec4<f32>({{ scaleBuffer }}[0]);
+  }
+  let base = hk * HEAD_DIM + d4 * 4u;
+  return vec4<f32>(
+    {{ scaleBuffer }}[base],
+    {{ scaleBuffer }}[base + 1u],
+    {{ scaleBuffer }}[base + 2u],
+    {{ scaleBuffer }}[base + 3u]
+  );
+}
+{%- endmacro %}
+{%- macro emit_quant_load4(format, kind, buffer, scaleBuffer) %}
+{{ emit_quant_scale4(kind, scaleBuffer) }}
+fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
+{%- if format == "int8" %}
+  return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
+{%- else %}
+  // Two elements cover this vec4: each carries two +8-biased nibbles, low first.
+  let rowBase = indexV4 - d4;
+  let lo = {{ buffer }}[rowBase + d4 * 2u];
+  let hi = {{ buffer }}[rowBase + d4 * 2u + 1u];
+  let nibbles = vec4<i32>(
+    i32(lo & 0xFu), i32((lo >> 4u) & 0xFu),
+    i32(hi & 0xFu), i32((hi >> 4u) & 0xFu)
+  );
+  let signed = nibbles - vec4<i32>(8);
+  return vec4<f32>(signed) * {{ kind }}scale4(d4, hk);
+{%- endif %}
+}
+{%- endmacro %}
+{{ emit_quant_load4("int8", "key", "key", "k_scale") }}
+{{ emit_quant_load4("int8", "value", "value", "v_scale") }}
+{% else %}
+fn load_key4(indexV4: u32) -> vec4<f32> {
+  return vec4<f32>(key[indexV4]);
+}
+fn load_value4(indexV4: u32) -> vec4<f32> {
+  return vec4<f32>(value[indexV4]);
+}
+{% endif %}
+{% if hasBias %}
+// Packed [Q; K; V] bias rows (token-independent). The Q bias folds into the
+// query row before the Q.K dots; the K bias adds a constant to every key score
+// that softmax cancels, so it is skipped; the V bias is token-independent and
+// is applied once in the merge pass after the final normalize.
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
+}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+  @builtin(subgroup_size) sgSize: u32{% endif %}
+) {
+{% if useSubgroups %}
+  // Subgroup tiles partition the fixed workgroup exactly. The advertised range
+  // is validated before dispatch; retain this uniform guard for implementations that
+  // choose an intermediate width at pipeline execution time.
+  if (sgSize == 0u || sgSize > WG || WG % sgSize != 0u) { return; }
+{% endif %}
+{% if splitQueries %}
+  let queryToken = wg.x / NUM_SPLITS;
+  let split = wg.x % NUM_SPLITS;
+{% else %}
+  let split = wg.x;
+{% endif %}
+  let h = wg.y;
+  let b = wg.z;
+  if (h >= Q_HEADS || split >= NUM_SPLITS{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}{% if source.layout == "layer_cache" %} || params.past_len >= CACHE_LEN{% endif %}) {
+    return;
+  }
+  let tid = lid.x;
+  let hKv = h / (Q_HEADS / KV_HEADS);
+{% if source.layout == "layer_cache" %}
+  let kvSeq = params.past_len + 1u;
+{% else %}
+  let cacheSeq = params.kvSeq;
+{% if cacheSeqlens %}
+  // Buffer-sharing caches retain their capacity in the physical BNSH stride;
+  // seqlens_k supplies the active end independently for each batch.
+  let kvSeq = min(cacheSeq, u32(seqlens_k[b]) + 1u);
+{% else %}
+  let kvSeq = cacheSeq;
+{% endif %}
+{% endif %}
+  // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
+{% if source.layout == "bsh" %}
+{% if splitQueries %}
+  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+{% endif %}
+  let kvBaseV4 = b * kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
+  let kvTokenStrideV4 = KV_HIDDEN_V4;
+{% elif source.layout == "layer_cache" %}
+  let qBaseV4 = h * HEAD_DIM_V4;
+  let kvBaseV4 = (LAYER * CACHE_LEN * KV_HEADS + hKv) * HEAD_DIM_V4;
+  let kvTokenStrideV4 = KV_HEADS * HEAD_DIM_V4;
+{% else %}
+{% if splitQueries %}
+  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
+{% else %}
+  let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
+{% endif %}
+  let kvBaseV4 = (b * KV_HEADS + hKv) * cacheSeq * HEAD_DIM_V4;
+  let kvTokenStrideV4 = HEAD_DIM_V4;
+{% endif %}
+  // Contiguous KV range owned by this split. Ceil division lets the last split
+  // absorb any remainder; empty ranges write identity partials and are ignored
+  // by the merge pass.
+{% if hasWindow %}
+  // Sliding window on the single decode query (absolute position
+  // kvSeq-1): it attends only the last `windowSize` keys, so split the
+  // contiguous [windowStart, kvSeq) range instead of the whole cache.
+  var windowStart: u32 = 0u;
+  if (kvSeq > params.windowSize) {
+    windowStart = kvSeq - params.windowSize;
+  }
+  let activeKeys = kvSeq - windowStart;
+  let keysPerSplit = (activeKeys + NUM_SPLITS - 1u) / NUM_SPLITS;
+  let splitStart = windowStart + split * keysPerSplit;
+{% else %}
+  let keysPerSplit = (kvSeq + NUM_SPLITS - 1u) / NUM_SPLITS;
+  let splitStart = split * keysPerSplit;
+{% endif %}
+  var splitEnd = splitStart + keysPerSplit;
+  if (splitEnd > kvSeq) {
+    splitEnd = kvSeq;
+  }
+{% set hasBias = hasBias is defined and hasBias %}
+  for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    var qv = vec4<f32>(query[qBaseV4 + d4]);
+{% if hasBias %}
+    qv = qv + load_bias4(h * HEAD_DIM, d4);
+{% endif %}
+    q_shared[d4] = qv;
+    running_out[d4] = vec4<f32>(0.0);
+  }
+  workgroupBarrier();
+  let scale = scale_value();
+  var runningMax = -FLT_MAX;
+  var runningDenom = 0.0;
+  var kjBase = splitStart;
+  loop {
+    if (kjBase >= splitEnd) {
+      break;
+    }
+    let kj = kjBase + tid;
+    var keyAllowed = kj < splitEnd;
+    let tileCount = min(WG, splitEnd - kjBase);
+    var score = -FLT_MAX;
+    var m = -FLT_MAX;
+    var dPart = 0.0;
+{% if coopQk %}
+    // Cooperative Q.K: one subgroup per key, lanes splitting HEAD_DIM_V4, then a hardware
+    // subgroupAdd — turns the per-thread HEAD_DIM_V4-long dependent dot chain into a few
+    // strided vec4 dots + one reduce. Uniform trip count keeps subgroupAdd in uniform flow.
+    let sgPerWg = WG / sgSize;
+    let qkRounds = (tileCount + sgPerWg - 1u) / sgPerWg;
+    let lane = tid % sgSize;
+    let sgInWg = tid / sgSize;
+    for (var rr: u32 = 0u; rr < qkRounds; rr = rr + 1u) {
+      let j = rr * sgPerWg + sgInWg;
+      var accS: f32 = 0.0;
+      if (j < tileCount) {
+        let kRowV4 = kvBaseV4 + (kjBase + j) * kvTokenStrideV4;
+        for (var d4: u32 = lane; d4 < HEAD_DIM_V4; d4 = d4 + sgSize) {
+          accS = accS + dot(q_shared[d4], load_key4(kRowV4 + d4{% if quantizedCache %}, d4, hKv{% endif %}));
+        }
+      }
+      let sj = subgroupAdd(accS);
+      if (lane == 0u && j < tileCount) {
+        sval_sh[j] = sj;
+      }
+    }
+    workgroupBarrier();
+    if (keyAllowed) {
+      score = sval_sh[tid] * scale;
+      m = score;
+      dPart = 1.0;
+    }
+{% else %}
+    if (keyAllowed) {
+      let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
+      var acc: f32 = 0.0;
+      for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+        acc = acc + dot(q_shared[d4], load_key4(kRowV4 + d4{% if quantizedCache %}, d4, hKv{% endif %}));
+      }
+      score = acc * scale;
+      m = score;
+      dPart = 1.0;
+    }
+{% endif %}
+{% if hasMask %}
+    if (keyAllowed) {
+{% if splitQueries %}
+      let maskQuery = queryToken;
+{% else %}
+      let maskQuery = 0u;
+{% endif %}
+      let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQuery * params.maskSeqStride + kj;
+{% if maskIsBool %}
+      // A rejected bool-mask key contributes no probability mass. The merge
+      // pass already maps a zero global denominator to an all-zero output row.
+      if (attn_mask[maskIndex] == 0u) {
+        keyAllowed = false;
+        score = -FLT_MAX;
+        dPart = 0.0;
+      }
+{% else %}
+      score = score + f32(attn_mask[maskIndex]);
+{% endif %}
+      m = score;
+    }
+{% endif %}
+    let tile = combine_partials(m, dPart, tid{% if useSubgroups %}, sgSize{% endif %});
+// Merge one key tile's online-softmax (maximum, denominator) partial into the
+// running state, then store the per-key probabilities consumed by V accumulation.
+    let newMax = max(runningMax, tile.x);
+    let correction = exp_shift(runningMax, newMax);
+    runningDenom = runningDenom * correction + tile.y * exp_shift(tile.x, newMax);
+    runningMax = newMax;
+    var prob = 0.0;
+    if (keyAllowed) {
+      prob = exp_shift(score, newMax);
+    }
+    probs[tid] = prob;
+    workgroupBarrier();
+{% if jSplitV %}
+    // j-split V accumulation: thread (jg, d4v) sums keys j == jg mod
+    // J_GROUPS for dim block d4v into a register, then the groups combine
+    // through shared memory so all lanes participate.
+    const J_GROUPS: u32 = {{ jGroups }}u;
+    let jg = tid / HEAD_DIM_V4;
+    let d4v = tid % HEAD_DIM_V4;
+    var vacc = vec4<f32>(0.0);
+    var jj = jg;
+    loop {
+      if (jj >= tileCount) { break; }
+      vacc = vacc + probs[jj] * load_value4(
+        kvBaseV4 + (kjBase + jj) * kvTokenStrideV4 + d4v{% if quantizedCache %},
+        d4v,
+        hKv{% endif %}
+      );
+      jj = jj + J_GROUPS;
+    }
+    vacc_sh[tid] = vacc;
+    workgroupBarrier();
+    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+      var a4 = running_out[d4] * correction;
+      for (var g: u32 = 0u; g < J_GROUPS; g = g + 1u) {
+        a4 = a4 + vacc_sh[g * HEAD_DIM_V4 + d4];
+      }
+      running_out[d4] = a4;
+    }
+    workgroupBarrier();
+{% else %}
+    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+      var vSum = vec4<f32>(0.0);
+      for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
+        vSum = vSum + probs[i] * load_value4(
+          kvBaseV4 + (kjBase + i) * kvTokenStrideV4 + d4{% if quantizedCache %},
+          d4,
+          hKv{% endif %}
+        );
+      }
+      running_out[d4] = running_out[d4] * correction + vSum;
+    }
+    workgroupBarrier();
+{% endif %}
+    kjBase = kjBase + WG;
+  }
+  // Emit un-normalized partials for (b, h, split): the merge pass divides.
+{% if splitQueries %}
+  let partialBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
+{% else %}
+  let partialBase = ((b * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
+{% endif %}
+  for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    partial_out[partialBase + d4] = running_out[d4];
+  }
+  if (tid == 0u) {
+{% if splitQueries %}
+    let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split;
+{% else %}
+    let mdBase = (b * Q_HEADS + h) * NUM_SPLITS + split;
+{% endif %}
+    // (max, denom) travel together to the merge, so they share one buffer as an
+    // interleaved vec2 rather than costing two bindings. Interleaved, not two
+    // halves, so the index needs no region size — and the merge reads both
+    // fields of a split in a single load.
+    partial_stats[mdBase] = vec2<f32>(runningMax, runningDenom);
+  }
+}

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

	@@ -0,0 +1,367 @@

+{% 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)
+// reduction. The running softmax state merges tiles with the stable online
+// max-and-denominator rule, and the headDim output accumulator is rescaled once
+// 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.
+{% endif %}
+{% if hasMask %}
+// Per-score mask uses broadcast strides. A stride of 0 broadcasts that axis;
+// rank-2 [q, k] masks set batch/head strides to 0. A rejected bool-mask key
+// contributes no softmax mass, and a fully masked row emits zero.
+{% endif %}
+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;
+// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
+// `m - m` finite so an empty lane / all--inf row contributes the exact
+// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
+// a zero final denominator according to their public semantics. Using -inf
+// here changes +inf-row behavior.
+const FLT_MAX: f32 = 3.4028234663852886e38;
+fn is_finite_f32(value: f32) -> bool {
+  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
+}
+// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
+// exactly at the row max. `x - x` on an infinite max is a legal fast-math
+// fold to 0, which would silently turn +inf rows finite — the explicit
+// equality test keeps the NaN propagation of the serial kernels.
+fn shifted_value(value: f32, maxValue: f32) -> f32 {
+  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
+  return select(value - maxValue, 0.0, equalFiniteMax);
+}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+var<workgroup> q_shared: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> running_out: array<vec4<f32>, HEAD_DIM_V4>;
+var<workgroup> probs: array<f32, WG>;
+// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
+// mNew = max(m1, m2)
+// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
+// Both the subgroup and portable barrier-tree engines return the same merged
+// pair to every invocation. Repeated merges require a workgroup barrier between
+// calls before their shared partial storage is reused.
+{% set combineSubgroups = combineSubgroups is defined and combineSubgroups %}
+{% if combineSubgroups %}
+// 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
+  // element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
+  // landed in d via exp(NaN); both survive the merge and are detected by the
+  // code after the reduction.
+  let sgD = subgroupAdd(d * exp_shift(m, sgM));
+  if (sgSize == WG) {
+    return vec2<f32>(sgM, sgD);
+  }
+  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);
+  }
+  workgroupBarrier();
+  return combinedMD;
+}
+{% else %}
+{% set mdStreamed = mdStreams is defined %}
+{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
+{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
+var<workgroup> partialM: array<f32, {{ mdExtent }}>;
+var<workgroup> partialD: array<f32, {{ mdExtent }}>;
+{% if mdStreamed %}
+// In-place fold of {{ mdStreams }} streams. 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;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+{% for s in range(mdStreams) %}
+      {
+        let slot = {{ s }}u * WG + lidx;
+        let m1 = partialM[slot];
+        let d1 = partialD[slot];
+        let m2 = partialM[slot + stride];
+        let d2 = partialD[slot + stride];
+        let mNew = max(m1, m2);
+        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+        partialM[slot] = mNew;
+      }
+{% endfor %}
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+}
+{% else %}
+fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
+  partialM[lidx] = m;
+  partialD[lidx] = d;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (lidx < stride) {
+      let m1 = partialM[lidx];
+      let d1 = partialD[lidx];
+      let m2 = partialM[lidx + stride];
+      let d2 = partialD[lidx + stride];
+      let mNew = max(m1, m2);
+      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
+      partialM[lidx] = mNew;
+    }
+    workgroupBarrier();
+    stride = stride / 2u;
+  }
+  let merged = vec2<f32>(partialM[0], partialD[0]);
+  // Trailing barrier so back-to-back calls cannot race a next call's partial
+  // stores against this call's reads of slot 0.
+  workgroupBarrier();
+  return merged;
+}
+{% endif %}
+{% endif %}
+// 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; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% if hasBias %}
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
+}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if combineSubgroups %},
+  @builtin(subgroup_size) sgSize: u32{% endif %}
+) {
+  let qi = wg.x;
+  let h = wg.y;
+  let b = wg.z;
+  if (qi >= params.qSeq || h >= {{ qHeads }}) {
+    return;
+  }
+  let tid = lid.x;
+  let hKv = h / ({{ qHeads }} / {{ kvHeads }});
+  // Row bases: Q/output share one row, K/V rows are kvBase + token * stride.
+  let qBaseV4 = (b * params.qSeq + qi) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+  let kvBaseV4 = b * params.kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
+  let kvTokenStrideV4 = KV_HIDDEN_V4;
+  // Cooperative vec4 Q-row load; init the output accumulator.
+{% set hasBias = hasBias is defined and hasBias %}
+  for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    var qv = vec4<f32>(query[qBaseV4 + d4]);
+{% if hasBias %}
+    qv = qv + load_bias4(h * HEAD_DIM, d4);
+{% endif %}
+    q_shared[d4] = qv;
+    running_out[d4] = vec4<f32>(0.0);
+  }
+  workgroupBarrier();
+  let scale = scale_value();
+  // Running (max, denom) kept redundantly in every thread's registers: both
+  // are derived only from workgroup-uniform tile merges, so every thread holds
+  // identical values without shared-memory round trips or extra barriers.
+  var runningMax = -FLT_MAX;
+  var runningDenom = 0.0;
+  // Causal upper bound: query qi attends only keys 0..qi, so stop after the tile
+  // containing qi and skip the unattended tail. Non-causal keeps the full kvSeq
+  // sweep.
+{% if hasCausal %}
+  var keyBoundV = select(params.kvSeq, min(params.kvSeq, qi + 1u), params.isCausal != 0u);
+{% else %}
+  var keyBoundV = params.kvSeq;
+{% endif %}
+  var keyFloor: u32 = 0u;
+{% if hasWindow %}
+  // Sliding window: query qi sits at absolute position p = kvSeq - qSeq
+  // + qi and attends keys [max(0, p+1-window), p+1) -- causal AND windowed.
+  let pastLen = select(0u, params.kvSeq - params.qSeq, params.kvSeq >= params.qSeq);
+  let qAbsP1 = pastLen + qi + 1u;
+  keyBoundV = min(keyBoundV, qAbsP1);
+  if (qAbsP1 > params.windowSize) { keyFloor = qAbsP1 - params.windowSize; }
+{% endif %}
+  let keyBound = keyBoundV;
+  // Start at the WG-aligned tile containing keyFloor (sub-floor lanes mask out).
+  let keyFloorAligned = (keyFloor / WG) * WG;
+  var kjBase: u32 = keyFloorAligned;
+  loop {
+    if (kjBase >= keyBound) {
+      break;
+    }
+    let kj = kjBase + tid;
+    // Thread-local (m, d) for this tile: one key per thread, identity
+    // (-FLT_MAX, 0) for out-of-range lanes on the tail tile.
+    var score = -FLT_MAX;
+    var m = -FLT_MAX;
+    var dPart = 0.0;
+    var keyAllowed = kj < keyBound{% if hasCausal %} && (params.isCausal == 0u || kj <= qi){% endif %}{% if hasWindow %} && kj >= keyFloor{% endif %};
+    if (keyAllowed) {
+      let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
+{% if hasMask %}
+      // Broadcast-strided mask address; a 0 stride collapses that axis (rank-2
+      // [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]));
+      }
+      score = acc * scale;
+{% if hasMask %}
+      score = score + f32(attn_mask[maskIndex]);
+{% endif %}
+{% endif %}
+      m = score;
+      dPart = select(0.0, 1.0, keyAllowed);
+    }
+    let tile = combine_partials(m, dPart, tid{% if combineSubgroups %}, sgSize{% endif %});
+    // Online merge of the tile into the running state (softmax-online rule).
+// Merge one key tile's online-softmax (maximum, denominator) partial into the
+// running state, then store the per-key probabilities consumed by V accumulation.
+    let newMax = max(runningMax, tile.x);
+    let correction = exp_shift(runningMax, newMax);
+    runningDenom = runningDenom * correction + tile.y * exp_shift(tile.x, newMax);
+    runningMax = newMax;
+    var prob = 0.0;
+    if (keyAllowed) {
+      prob = exp_shift(score, newMax);
+    }
+    probs[tid] = prob;
+    workgroupBarrier();
+    // running_out[d4] is owned by the same thread (tid ≡ d4 mod WG) across all
+    // tiles, so the rescale-and-accumulate below is race-free; vec4 V loads.
+    let tileCount = min(WG, keyBound - kjBase);
+    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+      var vSum = vec4<f32>(0.0);
+      for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
+        vSum = vSum + probs[i] * vec4<f32>(value[kvBaseV4 + (kjBase + i) * kvTokenStrideV4 + d4]);
+      }
+      running_out[d4] = running_out[d4] * correction + vSum;
+    }
+    // Orders this tile's probs reads before the next tile's probs writes.
+    workgroupBarrier();
+    kjBase = kjBase + WG;
+  }
+  // The output shares Q's layout and row base.
+  let hasMass = runningDenom > 0.0;
+  let invDenom = select(0.0, 1.0 / runningDenom, hasMass);
+  for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
+    var outValue = running_out[d4] * invDenom;
+{% if hasBias %}
+    // V bias is token-independent, so add it once after the softmax-weighted sum.
+    outValue = outValue + load_bias4(2u * {{ qHidden }}u + h * HEAD_DIM, d4)
+      * select(0.0, 1.0, hasMass);
+{% endif %}
+    output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue);
+  }
+}

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

	@@ -0,0 +1,554 @@

+{% 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" %}
+{% set Q_STRIDE = "QKV_STRIDE_V4" if sourceProfile == 1 else "Q_HIDDEN_V4" %}
+{% set QUERY = "qkv" if sourceProfile == 1 else "query" %}
+{% set KEY = "cache_keys" if sourceProfile == 1 else "key" %}
+{% set VALUE = "cache_values" if sourceProfile == 1 else "value" %}
+{% set OUTPUT = "attn_out" if sourceProfile == 1 else "output" %}
+{% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
+{% if batchNoSgReduction is not defined %}{% set batchNoSgReduction = false %}{% endif %}
+{% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
+{% if maskIsKeyKeep is not defined %}{% set maskIsKeyKeep = false %}{% endif %}
+{% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
+{% if hasSoftcap is not defined %}{% set hasSoftcap = false %}{% endif %}
+{% if hasHeadSink is not defined %}{% set hasHeadSink = false %}{% endif %}
+{% set Q_HIDDEN = qHidden | default(0) %}
+{% if quantCacheFormat is not defined %}{% set quantCacheFormat = "" %}{% endif %}
+{% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
+// A windowed cache binds a fixed CAPACITY but keeps only the most recent
+// min(total, capacity) rows resident. params.kvSeq then names the physical row
+// count, which is still the right batch stride but the wrong attention bound, so
+// the bounds read kvActive instead. Modes without seqlens use `params.kvSeq`
+// in both roles.
+{% set KVA = "kvActive" if useSeqlens else KVSEQ %}
+{% macro score_expr(part) %}{% if hasSoftcap %}params.softcap * tanh(clamp(({{ part }} * SCALE) / params.softcap, -30.0, 30.0)){% else %}{{ part }} * SCALE{% endif %}{% endmacro %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{% 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
+// workgroup storage; values widen to f32 when read. Other specializations stage
+// them as f32. Score and weighted-value accumulation remain in f32 throughout.
+{% set FUSED_ROTARY = fusedRotary is defined and fusedRotary %}
+{% set STAGE_MASK = hasMask and useSubgroups and stageMask is defined and stageMask %}
+{% set MASK_IS_INT = hasMask and (maskIsKeyKeep or maskIsBool) %}
+{% set MASK_TILE_TYPE = "u32" if MASK_IS_INT else "f32" %}
+{% set MASK_TILE_LOAD = "attn_mask[maskIndex]" if MASK_IS_INT else "f32(attn_mask[maskIndex])" %}
+{% set MASK_TILE_ZERO = "0u" if MASK_IS_INT else "0.0" %}
+{% set MASK_ELEMENT = "maskValue" if STAGE_MASK else "attn_mask[maskIndex]" %}
+{% set MASK_ADDITIVE = "maskValue" if STAGE_MASK else "f32(attn_mask[maskIndex])" %}
+{% set SLICE_COUNT = ((headDimV4 / LPQ) | int) %}
+{% set ROPE_LANE_XOR = ((LPQ / 2) | int) %}
+{% set QL = source.qLayout if source.qLayout is defined else source.layout %}
+{% set KL = source.kvLayout if source.kvLayout is defined else source.layout %}
+{% set RIGHT = causalRightAlign is defined and causalRightAlign %}
+// Tiled flash prefill attention with configurable-width query clusters for
+// token-major [batch, seq, heads*headDim] attention ops. Each workgroup covers
+// TILE_Q queries with one LPQ-lane cluster per query; K/V are staged in
+// workgroup memory once per TILE_K tile. Within a cluster, each lane holds q/o
+// register slices of HEAD_DIM/LPQ dimensions. Per-lane dot fragments are
+// combined with subgroup shuffles when available, or shared-memory reductions
+// on compatibility variants. Dot products, online softmax, and weighted-value
+// sums accumulate in f32. Optional masks are broadcast-strided; the
+// maskIsKeyKeep specialization supports 1=keep key-padding masks directly.
+//   hasCausal: params.isCausal masks keys after the query token (qSeq==kvSeq).
+//   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 bias (into qr) and the token-independent V bias (added
+//            once after the softmax-weighted sum) are applied. Each lane owns
+//            SLICE contiguous head-dim vec4s starting at lane8*SLICE.
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+{% if QL != "bhsd" %}
+const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;   // bsh token-major hidden stride
+{% endif %}
+{% if KL != "bhsd" %}
+const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
+{% endif %}
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const TILE_Q: u32 = {{ TILE_Q }}u;
+// Lanes per query cluster: 8 for headDim%32, 4 for headDim%16, 2 for headDim%8.
+const LPQ: u32 = {{ LPQ }}u;
+const SLICE: u32 = HEAD_DIM / (4u * LPQ);  // vec4s per lane slice
+{% if FUSED_ROTARY %}
+const HALF_DIM: u32 = HEAD_DIM / 2u;
+const HALF_LPQ: u32 = LPQ / 2u;
+{% endif %}
+const TILE_K: u32 = {{ TILE_K }}u;
+const WG: u32 = TILE_Q * LPQ;
+{% if MASK_IS_INT %}
+// Key-keep masks in contrib attention use a finite low logit for a rejected
+// key. Logical ONNX bool masks use the exclusion sentinel instead, so a fully
+// masked row has zero mass. Keeping one declaration shape for both mask modes
+// lets both mask modes share the score loop below.
+const NEG_INF: f32 = -3.4028234663852886e38;
+const MASK_NEG: f32 = {{ "-1e38" if maskIsKeyKeep else "-3.4028234663852886e38" }};
+{% else %}
+const NEG_INF: f32 = -3.4028234663852886e38;
+{% endif %}
+var<workgroup> k_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
+var<workgroup> v_tile: array<vec4<{{ ST }}>, TILE_K * (HEAD_DIM / 4u)>;
+{% if STAGE_MASK %}
+// Each LPQ cluster consumes one mask value per (query,key), so stage the
+// TILE_Q x TILE_K mask tile once instead of issuing LPQ duplicate global loads.
+var<workgroup> mask_tile: array<{{ MASK_TILE_TYPE }}, TILE_Q * TILE_K>;
+{% endif %}
+{% if not useSubgroups %}
+{% if batchNoSgReduction %}
+// No-subgroups cluster reduction scratch for a whole K tile. Staging every
+// lane's partial for every key lets the workgroup synchronize once per tile,
+// instead of twice per key. This path requires TILE_K*WG f32 scratch elements.
+var<workgroup> red: array<f32, TILE_K * WG>;
+{% else %}
+// One-key reduction path: reuse a single key's scratch across the tile.
+var<workgroup> red: array<f32, WG>;
+{% endif %}
+{% endif %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% if quantCacheFormat %}
+// A quantized cache is dequantized once per key into the staged tile, then read
+// by all TILE_Q queries in the workgroup. The unpack cost is amortized over the
+// tile height instead of paid once per (query, key).
+{% macro emit_quant_scale4(kind, scaleBuffer) %}
+fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
+  if (params.perChannel == 0u) {
+    return vec4<f32>({{ scaleBuffer }}[0]);
+  }
+  let base = hk * HEAD_DIM + d4 * 4u;
+  return vec4<f32>(
+    {{ scaleBuffer }}[base],
+    {{ scaleBuffer }}[base + 1u],
+    {{ scaleBuffer }}[base + 2u],
+    {{ scaleBuffer }}[base + 3u]
+  );
+}
+{%- endmacro %}
+{%- macro emit_quant_load4(format, kind, buffer, scaleBuffer) %}
+{{ emit_quant_scale4(kind, scaleBuffer) }}
+fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
+{%- if format == "int8" %}
+  return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
+{%- else %}
+  // Two elements cover this vec4: each carries two +8-biased nibbles, low first.
+  let rowBase = indexV4 - d4;
+  let lo = {{ buffer }}[rowBase + d4 * 2u];
+  let hi = {{ buffer }}[rowBase + d4 * 2u + 1u];
+  let nibbles = vec4<i32>(
+    i32(lo & 0xFu), i32((lo >> 4u) & 0xFu),
+    i32(hi & 0xFu), i32((hi >> 4u) & 0xFu)
+  );
+  let signed = nibbles - vec4<i32>(8);
+  return vec4<f32>(signed) * {{ kind }}scale4(d4, hk);
+{%- endif %}
+}
+{%- endmacro %}
+{{ emit_quant_load4(quantCacheFormat, "key", KEY, "k_scale") }}
+{{ emit_quant_load4(quantCacheFormat, "value", VALUE, "v_scale") }}
+{% endif %}
+{% if hasBias %}
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
+}
+{% endif %}
+@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 h = wg.y;
+  let b = wg.z;
+  let tid = lid.x;
+{% if useSubgroups %}
+  // The subgroupShuffleXor butterfly requires every LPQ-lane query cluster to
+  // fit in, and align with, one subgroup. Keep a uniform guard for defensive
+  // portability.
+  if (sgSize < LPQ || sgSize % LPQ != 0u) { return; }
+{% endif %}
+  let qSub = tid / LPQ;
+  let lane8 = tid % LPQ;
+  let qIdx = wg.x * TILE_Q + qSub;
+  let qValid = qIdx < {{ QSEQ }} && h < Q_HEADS;
+  let hKv = h / (Q_HEADS / KV_HEADS);
+  let SCALE = scale_value();
+  // Per-thread q slice + output accumulator (SLICE vec4s each) in registers.
+  let qClamped = min(qIdx, {{ QSEQ }} - 1u);
+{% if QL == "bhsd" %}
+  let qBase4 = ((b * Q_HEADS + h) * {{ QSEQ }} + qClamped) * HEAD_DIM_V4 + lane8 * SLICE;
+{% else %}
+  let qBase4 = (b * {{ QSEQ }} + qClamped) * {{ Q_STRIDE }} + h * HEAD_DIM_V4 + lane8 * SLICE;
+{% endif %}
+{% for c in range(SLICE_COUNT) %}
+  var qr{{ c }} = vec4<f32>({{ QUERY }}[qBase4 + {{ c }}u]);
+{% if hasBias %}
+  qr{{ c }} = qr{{ c }} + load_bias4(h * HEAD_DIM, lane8 * SLICE + {{ c }}u);
+{% endif %}
+  var o{{ c }} = vec4<f32>(0.0);
+{% endfor %}
+{% if FUSED_ROTARY %}
+  // The present-cache keys are already rotary transformed. Transform Q once in
+  // registers before the K loop, avoiding a separate full-size Q scratch pass.
+  // NeoX half-split pairs lanes separated by LPQ/2; subgroup variants exchange
+  // the paired register directly, while the portable path reloads that vec4.
+{% if not useSubgroups %}
+  let qHeadBase4 = qBase4 - lane8 * SLICE;
+  let pairedLane = (lane8 + HALF_LPQ) % LPQ;
+{% endif %}
+  let ropeChunk = (lane8 % HALF_LPQ) * SLICE;
+  let pastLenForRope = select(0u, {{ KVSEQ }} - {{ QSEQ }}, {{ KVSEQ }} >= {{ QSEQ }});
+  let ropePosition = pastLenForRope + qClamped;
+{% for c in range(SLICE_COUNT) %}
+{% if useSubgroups %}
+  let pairedQ{{ c }} = subgroupShuffleXor(qr{{ c }}, {{ ROPE_LANE_XOR }}u);
+{% else %}
+  let pairedQ{{ c }} = vec4<f32>({{ QUERY }}[qHeadBase4 + pairedLane * SLICE + {{ c }}u]);
+{% endif %}
+  let ropeBase{{ c }} = ropePosition * HALF_DIM + (ropeChunk + {{ c }}u) * 4u;
+  let ropeCos{{ c }} = vec4<f32>(
+    f32(cos_cache[ropeBase{{ c }}]),
+    f32(cos_cache[ropeBase{{ c }} + 1u]),
+    f32(cos_cache[ropeBase{{ c }} + 2u]),
+    f32(cos_cache[ropeBase{{ c }} + 3u])
+  );
+  let ropeSin{{ c }} = vec4<f32>(
+    f32(sin_cache[ropeBase{{ c }}]),
+    f32(sin_cache[ropeBase{{ c }} + 1u]),
+    f32(sin_cache[ropeBase{{ c }} + 2u]),
+    f32(sin_cache[ropeBase{{ c }} + 3u])
+  );
+  qr{{ c }} = select(
+    qr{{ c }} * ropeCos{{ c }} + pairedQ{{ c }} * ropeSin{{ c }},
+    qr{{ c }} * ropeCos{{ c }} - pairedQ{{ c }} * ropeSin{{ c }},
+    lane8 < HALF_LPQ
+  );
+{% endfor %}
+{% endif %}
+  var m: f32 = NEG_INF;
+  var l: f32 = 0.0;
+{% if useSeqlens %}
+  // Resident rows of a windowed cache: the survivors were shifted down to [0, kvActive),
+  // so query/key DISTANCE is unchanged and every bound below reads as if the cache were
+  // exactly kvActive long. Rotary is excluded from this path (it would need the absolute
+  // position, not the cache-relative one), so pastLenForRope keeps the physical length.
+  let kvActive = min({{ KVSEQ }}, u32(seqlens_k[b]) + 1u);
+{% endif %}
+  // Causal ceiling per query; the key loop runs over the workgroup's union range
+  // (uniform trip count), masking out-of-range (query, key) pairs.
+{% 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
+  // windowed, the floor p+1-windowSize. The key loop runs over the workgroup's
+  // union range [wgStart, wgEnd); out-of-range (query,key) pairs mask to NEG_INF.
+  // 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 RIGHT or hasWindow %}
+  let pastLen = select(0u, {{ KVA }} - {{ QSEQ }}, {{ KVA }} >= {{ QSEQ }});
+{% endif %}
+{% if hasCausal %}
+{% if RIGHT %}
+  var maxKjV = select({{ KVA }}, min(pastLen + qIdx + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
+{% else %}
+  var maxKjV = select({{ KVA }}, min(qIdx + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
+{% endif %}
+{% else %}
+  var maxKjV = {{ KVA }};
+{% endif %}
+  var minKjV: u32 = 0u;
+  let lastQ = min(wg.x * TILE_Q + TILE_Q - 1u, {{ QSEQ }} - 1u);
+{% if hasCausal %}
+{% if RIGHT %}
+  var wgEndV = select({{ KVA }}, min(pastLen + lastQ + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
+{% else %}
+  var wgEndV = select({{ KVA }}, min(lastQ + 1u, {{ KVA }}), {{ IS_CAUSAL }} != 0u);
+{% endif %}
+{% else %}
+  var wgEndV = {{ KVA }};
+{% endif %}
+  var wgStartV: u32 = 0u;
+{% if hasWindow %}
+  let qAbsP1 = pastLen + qIdx + 1u;
+  maxKjV = min(maxKjV, qAbsP1);
+  if (qAbsP1 > params.windowSize) { minKjV = qAbsP1 - params.windowSize; }
+  let lastQAbsP1 = pastLen + lastQ + 1u;
+  wgEndV = min(wgEndV, lastQAbsP1);
+  let firstQAbsP1 = pastLen + wg.x * TILE_Q + 1u;
+  if (firstQAbsP1 > params.windowSize) { wgStartV = firstQAbsP1 - params.windowSize; }
+{% endif %}
+  let maxKj = maxKjV;
+  let minKj = minKjV;
+  let wgEnd = wgEndV;
+  let wgStart = wgStartV;
+{% else %}
+  let maxKj = {{ KVA }};
+  let minKj: u32 = 0u;
+  let wgEnd = {{ KVA }};
+  let wgStart: u32 = 0u;
+{% endif %}
+{% if KL == "bhsd" %}
+{% if quantCacheFormat == "int4" %}
+  let kvRowStride4 = HEAD_DIM_V4 * 2u;
+{% else %}
+  let kvRowStride4 = HEAD_DIM_V4;
+{% endif %}
+  let kvBatch4 = (b * KV_HEADS + hKv) * {{ KVSEQ }} * kvRowStride4;
+{% else %}
+  let kvBatch4 = b * {{ KVSEQ }} * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
+{% endif %}
+  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;
+      let d4 = i % HEAD_DIM_V4;
+      let kj = kStart + slot;
+      if (kj < wgEnd) {
+{% if KL == "bhsd" %}
+        let base4 = kvBatch4 + kj * kvRowStride4 + d4;
+{% else %}
+        let base4 = kvBatch4 + kj * KV_HIDDEN_V4 + d4;
+{% endif %}
+{% if quantCacheFormat %}
+        k_tile[i] = vec4<{{ ST }}>(load_key4(base4, d4, hKv));
+        v_tile[i] = vec4<{{ ST }}>(load_value4(base4, d4, hKv));
+{% else %}
+        k_tile[i] = vec4<{{ ST }}>({{ KEY }}[base4]);
+        v_tile[i] = vec4<{{ ST }}>({{ VALUE }}[base4]);
+{% endif %}
+      } else {
+        k_tile[i] = vec4<{{ ST }}>(0.0);
+        v_tile[i] = vec4<{{ ST }}>(0.0);
+      }
+    }
+{% if STAGE_MASK %}
+    // The K/V load barrier also publishes this compact mask tile.
+    for (var i: u32 = tid; i < TILE_Q * TILE_K; i = i + WG) {
+      let qSlot = i / TILE_K;
+      let kSlot = i % TILE_K;
+      let maskQ = min(wg.x * TILE_Q + qSlot, {{ QSEQ }} - 1u);
+      let maskK = kStart + kSlot;
+      if (maskK < wgEnd) {
+        let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQ * params.maskSeqStride + maskK;
+        mask_tile[i] = {{ MASK_TILE_LOAD }};
+      } else {
+        mask_tile[i] = {{ MASK_TILE_ZERO }};
+      }
+    }
+{% endif %}
+    workgroupBarrier();
+    // TILE_K 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>;
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+      s[kk] = NEG_INF;
+      var part: f32 = 0.0;
+      let kb = kk * HEAD_DIM_V4 + lane8 * SLICE;
+{% for c in range(SLICE_COUNT) %}
+      part = part + dot(qr{{ c }}, vec4<f32>(k_tile[kb + {{ c }}u]));
+{% endfor %}
+      red[kk * WG + tid] = part;
+    }
+    workgroupBarrier();
+    // Distribute the tile's key reductions across the LPQ lanes, then publish
+    // each score at the cluster base. This replaces LPQ identical gathers and
+    // avoids making lane zero serially reduce the whole tile. The extra barrier
+    // is still paid once for the whole K tile, not once per key.
+    let cbase = tid - lane8;
+    for (var kk: u32 = lane8; kk < TILE_K; kk = kk + LPQ) {
+{% for j in range(LPQ) %}
+      let r{{ j }} = red[kk * WG + cbase + {{ j }}u];
+{% endfor %}
+{% if LPQ == 8 %}
+      let part = ((r0 + r1) + (r2 + r3)) + ((r4 + r5) + (r6 + r7));
+{% elif LPQ == 4 %}
+      let part = (r0 + r1) + (r2 + r3);
+{% else %}
+      let part = r0 + r1;
+{% endif %}
+      red[kk * WG + cbase] = part;
+    }
+    workgroupBarrier();
+    // The barrier at the start of the next K tile protects this scratch before
+    // it is overwritten.
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+      let kj = kStart + kk;
+      let part = red[kk * WG + tid - lane8];
+      if (kj >= minKj && kj < maxKj) {
+{% if hasMask %}
+        let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
+{% if maskIsBool %}
+        if (attn_mask[maskIndex] != 0u) {
+          s[kk] = {{ score_expr("part") }};
+        } else {
+          s[kk] = MASK_NEG;
+        }
+{% else %}
+        s[kk] = {{ score_expr("part") }} + f32(attn_mask[maskIndex]);
+{% endif %}
+{% else %}
+        s[kk] = {{ score_expr("part") }};
+{% endif %}
+      }
+    }
+{% else %}
+    var s: array<f32, TILE_K>;
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+      let kj = kStart + kk;
+      s[kk] = NEG_INF;
+      var part: f32 = 0.0;
+      let kb = kk * HEAD_DIM_V4 + lane8 * SLICE;
+{% for c in range(SLICE_COUNT) %}
+      part = part + dot(qr{{ c }}, vec4<f32>(k_tile[kb + {{ c }}u]));
+{% endfor %}
+{% if useSubgroups %}
+      // log2(LPQ) butterfly all-reduce within the LPQ-lane cluster (each lane ends with the full dot).
+      part = part + subgroupShuffleXor(part, 1u);
+{% for shift in [2, 4] %}
+{% if LPQ > shift %}
+      part = part + subgroupShuffleXor(part, {{ shift }}u);
+{% endif %}
+{% endfor %}
+{% else %}
+      red[tid] = part;
+      workgroupBarrier();
+      let cbase = tid - lane8;
+      part = 0.0;
+      for (var j: u32 = 0u; j < LPQ; j = j + 1u) {
+        part = part + red[cbase + j];
+      }
+      workgroupBarrier();
+{% endif %}
+      if (kj >= minKj && kj < maxKj) {
+{% if hasMask %}
+{% if STAGE_MASK %}
+        let maskValue = mask_tile[qSub * TILE_K + kk];
+{% else %}
+        // Broadcast-strided mask address (a 0 stride collapses that axis; rank-2
+        // [q, k] masks set batch/head strides to 0). qClamped keeps the seq index
+        // in-bounds for padding queries in the last tile (their output is dropped).
+        let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + qClamped * params.maskSeqStride + kj;
+{% endif %}
+{% if maskIsKeyKeep %}
+        // A broadcast key mask uses 1 for a retained key and 0 for padding.
+        s[kk] = {{ score_expr("part") }} + (1.0 - f32({{ MASK_ELEMENT }})) * MASK_NEG;
+{% elif maskIsBool %}
+        // Logical bool: a rejected key contributes no softmax mass. Leaving
+        // the initialized NEG_INF sentinel in place makes a fully masked row
+        // land on the zero-denominator output guard below.
+        if ({{ MASK_ELEMENT }} != 0u) {
+          s[kk] = {{ score_expr("part") }};
+        } else {
+          s[kk] = MASK_NEG;
+        }
+{% else %}
+        s[kk] = {{ score_expr("part") }} + {{ MASK_ADDITIVE }};
+{% endif %}
+{% else %}
+        s[kk] = {{ score_expr("part") }};
+{% endif %}
+      }
+    }
+{% endif %}
+    // 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) %}
+    {
+      var acc = o{{ c }} * corr;
+      if (tileInterior) {
+        for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+          acc = acc + s[kk] * vec4<f32>(v_tile[kk * HEAD_DIM_V4 + lane8 * SLICE + {{ c }}u]);
+        }
+      } else {
+        for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+          acc = acc + s[kk] * select(vec4<f32>(), vec4<f32>(v_tile[kk * HEAD_DIM_V4 + lane8 * SLICE + {{ c }}u]), s[kk] != 0.0);
+        }
+      }
+      o{{ c }} = acc;
+    }
+{% endfor %}
+    kStart = kStart + TILE_K;
+  }
+  if (qValid) {
+{% if QL == "bhsd" %}
+    let outBase4 = ((b * Q_HEADS + h) * {{ QSEQ }} + qIdx) * HEAD_DIM_V4 + lane8 * SLICE;
+{% else %}
+    let outBase4 = (b * {{ QSEQ }} + qIdx) * Q_HIDDEN_V4 + h * HEAD_DIM_V4 + lane8 * SLICE;
+{% endif %}
+{% if hasHeadSink %}
+    // The head sink is a learned logit that competes with the keys but carries
+    // no value, so it enters the denominator only and the weighted sum above is
+    // untouched. Renormalizing against max(m, sink) keeps the exponentials in
+    // range when the sink dominates a fully-masked row.
+    let sink = f32(head_sink[h]);
+    let finalM = max(m, sink);
+    let accScale = exp(m - finalM);
+    let inv = accScale / (exp(sink - finalM) + l * accScale);
+{% else %}
+    // l == 0 means this query had no probability-bearing key: either its
+    // causal/window range is empty or its logical bool mask rejects every key.
+    // Emit 0 rather than 0/0. Other attention paths enforce the same empty-row
+    // contract by selecting on positive global mass.
+    let inv = select(0.0, 1.0 / l, l > 0.0);
+{% endif %}
+{% macro attention_value(c) %}o{{ c }} * inv{% if hasBias %} + load_bias4(2u * {{ Q_HIDDEN }}u + h * HEAD_DIM, lane8 * SLICE + {{ c }}u){% endif %}{% endmacro %}
+{% for c in range(SLICE_COUNT) %}
+    {{ OUTPUT }}[outBase4 + {{ c }}u] = vec4<{{ scalar }}>({{ attention_value(c) }});
+{% endfor %}
+  }
+}

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

	@@ -0,0 +1,246 @@

+// Register-resident flash prefill uses one 32-lane subgroup per workgroup. Each
+// lane owns a query row and keeps its Q slice and f32 output in registers. Lanes
+// cooperatively load K/V and broadcast them with subgroupShuffle, so each query
+// computes q·k and p·v without cross-lane reductions or workgroup storage.
+//
+// Batch, BHSD layout, grouped heads, and causal key ceilings share one uniform
+// key-loop bound so subgroup operations stay reconverged. Packed Q/K/V bias
+// applies Q before attention and V after normalization; key-independent K bias
+// 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 %}
+{{ env.wgsl.resourceDeclarations }}
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const Q_HEADS:     u32 = {{ qNumHeads }}u;
+const KV_HEADS:    u32 = {{ kvNumHeads }}u;
+const K_STEP:      u32 = {{ kStep }}u;
+const Q_STEP:      u32 = {{ Q_STEP }}u;
+{% if hasBias %}
+const HEAD_DIM:    u32 = {{ headDim }}u;
+{% endif %}
+const Q_HIDDEN_V4:  u32 = Q_HEADS * HEAD_DIM_V4;
+const KV_HIDDEN_V4: u32 = KV_HEADS * HEAD_DIM_V4;
+const NEG_INF:     f32 = -3.4028234663852886e38;
+{% if not USE_SUBGROUPS %}
+// No-subgroup backend: the same one-query-per-lane algorithm broadcasts a
+// 32-key K/V tile through workgroup memory. The portable fallback can assign
+// 64 query rows to one workgroup, so two SIMD-widths share each staged tile and
+// halve K/V traffic without subgroup operations. At headDim=256 the two f16
+// tiles consume exactly the portable 32 KiB workgroup-storage limit.
+var<workgroup> keyTile: array<vec4<{{ ST }}>, HEAD_DIM_V4 * K_STEP>;
+var<workgroup> valueTile: array<vec4<{{ ST }}>, HEAD_DIM_V4 * K_STEP>;
+{% endif %}
+{% set ATTN_SCALE_DIM = headDim ~ "u" %}
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% if hasBias %}
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
+}
+{% endif %}
+@compute @workgroup_size({{ Q_STEP }}, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if USE_SUBGROUPS %},
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+{% endif %}
+) {
+{% if USE_SUBGROUPS %}
+  if (sg_size != 32u) { return; }
+{% else %}
+  let lane = lid.x;
+{% endif %}
+  let h = wg.y;
+  let b = wg.z;
+  let qi = wg.x * Q_STEP + lid.x;
+  let valid_q = qi < params.qSeq;
+  let hKv = h / (Q_HEADS / KV_HEADS);
+  let SCALE = scale_value();
+  // Q row. Padding lanes (qi >= qSeq) read query 0-clamped and never store, so they
+  // stay NaN-free while keeping the subgroup full for the cooperative K/V broadcast.
+  let qClamped = min(qi, params.qSeq - 1u);
+  let qBaseV4 = (b * params.qSeq + qClamped) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+  var q_tile: array<vec4<{{ ST }}>, HEAD_DIM_V4>;
+  var o_tile: array<vec4<f32>, HEAD_DIM_V4>;
+  for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+    if (valid_q) {
+      q_tile[d4] = query[qBaseV4 + d4];
+{% if hasBias %}
+      q_tile[d4] = vec4<{{ ST }}>(vec4<f32>(q_tile[d4]) + load_bias4(h * HEAD_DIM, d4));
+{% endif %}
+    } else {
+      q_tile[d4] = vec4<{{ ST }}>(0.0);
+    }
+    o_tile[d4] = vec4<f32>(0.0);
+  }
+  var previous_max: f32 = NEG_INF;
+  var previous_denom: f32 = 0.0;
+  // Causal key ceiling. The key-loop bound (kvEnd) uses the workgroup's LAST query
+  // so every lane shares a uniform trip count (subgroup ops stay reconverged);
+  // each lane masks its own keys past myMaxKj to NEG_INF.
+{% if CAUSAL %}
+  let lastQ = min(wg.x * Q_STEP + Q_STEP - 1u, params.qSeq - 1u);
+  let kvEnd = select(params.kvSeq, min(lastQ + 1u, params.kvSeq), params.isCausal != 0u);
+  let myMaxKj = select(params.kvSeq, min(qi + 1u, params.kvSeq), params.isCausal != 0u);
+{% else %}
+  let kvEnd = params.kvSeq;
+  let myMaxKj = params.kvSeq;
+{% endif %}
+  let kvBase = b * params.kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
+  let kvKeyStride = KV_HIDDEN_V4;
+  for (var k_start: u32 = 0u; k_start < kvEnd; k_start = k_start + K_STEP) {
+{% if not USE_SUBGROUPS %}
+    // Cooperatively stage K and V once, then all query lanes read the complete
+    // tile directly. A trailing barrier below protects the next tile overwrite.
+    for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+      for (var ki = lane; ki < K_STEP; ki += Q_STEP) {
+        let tileOffset = d4 * K_STEP + ki;
+        var kLocal = vec4<{{ ST }}>(0.0);
+        var vLocal = vec4<{{ ST }}>(0.0);
+        if (k_start + ki < params.kvSeq) {
+          let sourceOffset = kvBase + (k_start + ki) * kvKeyStride + d4;
+          kLocal = key[sourceOffset];
+          vLocal = value[sourceOffset];
+        }
+        keyTile[tileOffset] = kLocal;
+        valueTile[tileOffset] = vLocal;
+      }
+    }
+    workgroupBarrier();
+{% endif %}
+{% for g in range(qkGroups) %}
+    var qk{{ g }}: vec4<f32> = vec4<f32>(0.0);
+{% endfor %}
+    for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+      let q_own = q_tile[d4];
+{% if USE_SUBGROUPS %}
+      var k_local0: vec4<{{ ST }}> = vec4<{{ ST }}>(0.0);
+{% if kStep == 64 %}
+      var k_local1: vec4<{{ ST }}> = vec4<{{ ST }}>(0.0);
+{% endif %}
+      if (k_start + sg_id < params.kvSeq) {
+        k_local0 = key[kvBase + (k_start + sg_id) * kvKeyStride + d4];
+      }
+{% if kStep == 64 %}
+      if (k_start + sg_id + 32u < params.kvSeq) {
+        k_local1 = key[kvBase + (k_start + sg_id + 32u) * kvKeyStride + d4];
+      }
+{% endif %}
+{% endif %}
+{% for g in range(qkGroups) %}
+{% for lane in range(4) %}
+{% if USE_SUBGROUPS %}
+{% if g < 8 %}
+      qk{{ g }}.{{ components[lane] }} = qk{{ g }}.{{ components[lane] }} + f32(dot(q_own, subgroupShuffle(k_local0, {{ g * 4 + lane }}u)));
+{% else %}
+      qk{{ g }}.{{ components[lane] }} = qk{{ g }}.{{ components[lane] }} + f32(dot(q_own, subgroupShuffle(k_local1, {{ (g - 8) * 4 + lane }}u)));
+{% endif %}
+{% else %}
+      qk{{ g }}.{{ components[lane] }} = qk{{ g }}.{{ components[lane] }} + f32(dot(q_own, keyTile[d4 * K_STEP + {{ g * 4 + lane }}u]));
+{% endif %}
+{% endfor %}
+{% endfor %}
+    }
+    // Scale + range/causal mask (a masked score is the NEG_INF sentinel, which the
+    // select below turns into a 0 weight — never an exp(NEG_INF-NEG_INF)=1 leak).
+{% for g in range(qkGroups) %}
+    qk{{ g }} = qk{{ g }} * vec4<f32>(SCALE);
+{% for lane in range(4) %}
+    if (k_start + {{ g * 4 + lane }}u >= myMaxKj) { qk{{ g }}.{{ components[lane] }} = NEG_INF; }
+{% endfor %}
+{% endfor %}
+    // Per-lane online softmax (this lane's query is independent — no barrier).
+    var local_max: f32 = NEG_INF;
+{% for g in range(qkGroups) %}
+    local_max = max(local_max, max(max(qk{{ g }}.x, qk{{ g }}.y), max(qk{{ g }}.z, qk{{ g }}.w)));
+{% endfor %}
+    let new_max = max(previous_max, local_max);
+{% for g in range(qkGroups) %}
+    qk{{ g }} = select(vec4<f32>(0.0), exp(qk{{ g }} - vec4<f32>(new_max)), qk{{ g }} != vec4<f32>(NEG_INF));
+{% endfor %}
+    var tile_sum: f32 = 0.0;
+{% for g in range(qkGroups) %}
+    tile_sum = tile_sum + qk{{ g }}.x + qk{{ g }}.y + qk{{ g }}.z + qk{{ g }}.w;
+{% endfor %}
+    let dleft = previous_denom * exp(previous_max - new_max);
+    let denom = max(dleft + tile_sum, 0.0000001);
+    let o_ratio = dleft / denom;
+{% for g in range(qkGroups) %}
+    qk{{ g }} = qk{{ g }} / vec4<f32>(denom);
+{% endfor %}
+    previous_max = new_max;
+    previous_denom = denom;
+    for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+{% if USE_SUBGROUPS %}
+      var v_local0: vec4<{{ ST }}> = vec4<{{ ST }}>(0.0);
+{% if kStep == 64 %}
+      var v_local1: vec4<{{ ST }}> = vec4<{{ ST }}>(0.0);
+{% endif %}
+      if (k_start + sg_id < params.kvSeq) {
+        v_local0 = value[kvBase + (k_start + sg_id) * kvKeyStride + d4];
+      }
+{% if kStep == 64 %}
+      if (k_start + sg_id + 32u < params.kvSeq) {
+        v_local1 = value[kvBase + (k_start + sg_id + 32u) * kvKeyStride + d4];
+      }
+{% endif %}
+{% endif %}
+      var acc: vec4<f32> = vec4<f32>(0.0);
+{% for g in range(qkGroups) %}
+{% for lane in range(4) %}
+{% if USE_SUBGROUPS %}
+{% if g < 8 %}
+      acc = acc + vec4<f32>(subgroupShuffle(v_local0, {{ g * 4 + lane }}u)) * qk{{ g }}.{{ components[lane] }};
+{% else %}
+      acc = acc + vec4<f32>(subgroupShuffle(v_local1, {{ (g - 8) * 4 + lane }}u)) * qk{{ g }}.{{ components[lane] }};
+{% endif %}
+{% else %}
+      acc = acc + vec4<f32>(valueTile[d4 * K_STEP + {{ g * 4 + lane }}u]) * qk{{ g }}.{{ components[lane] }};
+{% endif %}
+{% endfor %}
+{% endfor %}
+      o_tile[d4] = o_tile[d4] * vec4<f32>(o_ratio) + acc;
+    }
+{% if not USE_SUBGROUPS %}
+    workgroupBarrier();
+{% endif %}
+  }
+  if (valid_q) {
+    let outBaseV4 = (b * params.qSeq + qi) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+    for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
+{% if hasBias %}
+      output[outBaseV4 + d4] = vec4<{{ scalar }}>(o_tile[d4] + load_bias4(2u * {{ qHidden }}u + h * HEAD_DIM, d4));
+{% else %}
+      output[outBaseV4 + d4] = vec4<{{ scalar }}>(o_tile[d4]);
+{% endif %}
+    }
+  }
+}

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

	@@ -0,0 +1,205 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set MASK_BATCH = "batch * params.maskBatchStride + " if source.layout == "bsh" else "" %}
+// Online-softmax attention fallback with no feature requirements: one
+// workgroup per (batch, head, query token) walks the keys serially; the
+// workgroup cooperates on each q·k dot (tree reduction) and on the running
+// V accumulator, with the online rescale applied per key. It is the portable
+// path for any workgroup size and device tier.
+// Layout: rank-3 token-major [batch, seq, heads * headDim].
+{% if headsFromParams %}
+// Q/KV head counts are uniforms; the GQA mapping is h / (qHeads / kvHeads).
+{% endif %}
+{% if hasKeyLimit %}
+// params.keyLimit caps the key range when nonzero.
+{% endif %}
+{% if scaleFallbackRsqrt %}
+// An omitted scale uses 1/sqrt(headDim); explicit zero remains zero through
+// the scaleIsExplicitZero specialization.
+{% endif %}
+{% if hasBias %}
+// Packed [Q; K; V] bias rows are applied during the serial key walk.
+{% endif %}
+{% if hasMask %}
+// Additive per-score attention bias with broadcast strides: a stride of 0
+// collapses that axis (batch and/or head broadcast).
+{% endif %}
+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") %}
+const WG: u32 = {{ workgroupSize }}u;
+var<workgroup> partial: array<f32, WG>;
+var<workgroup> running_max: f32;
+var<workgroup> running_denom: f32;
+var<workgroup> running_out: array<f32, HEAD_DIM>;
+var<workgroup> previous_scale: f32;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
+fn {{ name }}(value: f32, tid: u32) -> f32 {
+  {{ buffer }}[tid] = value;
+  workgroupBarrier();
+  // Ceil-halving keeps every lane when the workgroup size is not a power of
+  // two. For even n this matches the power-of-two tree order; for odd n, lanes
+  // [0, n-half) fold the upper tail while the middle lane carries forward.
+  var n: u32 = {{ wg }};
+  loop {
+    let half = (n + 1u) / 2u;
+    if (tid < n - half) {
+{% if mode == "max" %}
+      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
+{% else %}
+      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
+{% endif %}
+    }
+    workgroupBarrier();
+    n = half;
+    if (n == 1u) {
+      break;
+    }
+  }
+  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
+  // slot 0 here, so the next call's first store must not run until all lanes have read it.
+  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
+  let reduced = {{ buffer }}[0];
+{% if trailingBarrier %}
+  workgroupBarrier();
+{% endif %}
+  return reduced;
+}
+{% endmacro %}
+{{ 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(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let h = wg.y;
+  let batch = wg.z;
+  if (h >= {{ qHeads }}) {
+    return;
+  }
+  let tid = lid.x;
+  let kv_head = h / ({{ qHeads }} / {{ kvHeads }});
+  // Grid-stride over query tokens: the dispatch clamps x to maxComputeWorkgroupsPerDimension, so for
+  // qSeq > maxComputeWorkgroupsPerDimension each workgroup serially handles queries wg.x, wg.x + nwg.x, …
+  // Every query's online-softmax state is workgroup-local and re-initialized
+  // each iteration, so iterations are independent; the trailing workgroupBarrier
+  // keeps a query's reads of running_denom ahead of the next query's writes.
+  // Reduces to exactly one query per workgroup (bit-identical) when nwg.x == qSeq.
+  for (var query_token = wg.x; query_token < params.qSeq; query_token = query_token + nwg.x) {
+  // Row bases: Q/output share one row, K/V rows are kvBase + token * stride.
+  let qBase = (batch * params.qSeq + query_token) * Q_HIDDEN + h * HEAD_DIM;
+  let kvBase = batch * params.kvSeq * KV_HIDDEN + kv_head * HEAD_DIM;
+  let kvTokenStride = KV_HIDDEN;
+  if (tid == 0u) {
+    running_max = -3.4028234663852886e38;
+    running_denom = 0.0;
+  }
+  for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
+    running_out[d] = 0.0;
+  }
+  workgroupBarrier();
+  var maxKj = params.kvSeq;
+  var minKj: u32 = 0u;
+{% if hasCausal %}
+  maxKj = min(maxKj, select(params.kvSeq, query_token + 1u, params.isCausal != 0u));
+{% endif %}
+{% if hasKeyLimit %}
+  maxKj = min(maxKj, select(params.kvSeq, params.keyLimit, params.keyLimit > 0u));
+{% endif %}
+{% if hasWindow %}
+  // local_window_size: the query at relative index
+  // query_token sits at absolute position p = kvSeq - qSeq + query_token, so it
+  // attends keys [max(0, p + 1 - window), p + 1) -- causal AND left-windowed.
+  let pastLen = select(0u, params.kvSeq - params.qSeq, params.kvSeq >= params.qSeq);
+  let qAbsP1 = pastLen + query_token + 1u;
+  maxKj = min(maxKj, qAbsP1);
+  if (qAbsP1 > params.windowSize) {
+    minKj = qAbsP1 - params.windowSize;
+  }
+{% endif %}
+  for (var key_token: u32 = minKj; key_token < maxKj; key_token = key_token + 1u) {
+    let kRow = kvBase + key_token * kvTokenStride;
+    var partial_dot = 0.0;
+    for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
+      var q_value = f32(query[qBase + d]);
+      var k_value = f32(key[kRow + d]);
+{% if hasBias %}
+      let channel = h * HEAD_DIM + d;
+      q_value = q_value + f32(bias[channel]);
+      k_value = k_value + f32(bias[Q_HIDDEN + channel]);
+{% endif %}
+      partial_dot = partial_dot + q_value * k_value;
+    }
+    // reduce_sum returns the same partial[0] to every lane, so `score` is already
+    // workgroup-uniform; tid==0 advances the running max/denom/scale (broadcast
+    // through shared memory by the barrier below).
+{% if hasMask %}
+    let maskIndex = {{ MASK_BATCH }}h * params.maskHeadStride + query_token * params.maskSeqStride + key_token;
+    let score = reduce_sum(partial_dot, tid) * scale_value() + f32(attn_mask[maskIndex]);
+{% else %}
+    let score = reduce_sum(partial_dot, tid) * scale_value();
+{% endif %}
+    if (tid == 0u) {
+      let next_max = max(running_max, score);
+      let old_factor = exp(running_max - next_max);
+      let new_factor = exp(score - next_max);
+      previous_scale = old_factor;
+      running_max = next_max;
+      running_denom = running_denom * old_factor + new_factor;
+    }
+    workgroupBarrier();
+    let probability_numerator = exp(score - running_max);
+    for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
+      var v_value = f32(value[kRow + d]);
+{% if hasBias %}
+      v_value = v_value + f32(bias[2u * Q_HIDDEN + h * HEAD_DIM + d]);
+{% endif %}
+      running_out[d] = running_out[d] * previous_scale + probability_numerator * v_value;
+    }
+    workgroupBarrier();
+  }
+  let inv_denom = 1.0 / running_denom;
+  for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
+    output[qBase + d] = {{ outputScalar }}(running_out[d] * inv_denom);
+  }
+  // Separate this query's reads of the shared running state from the next
+  // iteration's re-initialization (grid-stride; no-op when the loop runs once).
+  workgroupBarrier();
+  }
+}

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,1565 @@

+{
+  "op": "com.microsoft.GroupQueryAttention",
+  "tunableSpace": {
+    "PREFILL_QUERY_TILE": [16, 32],
+    "PREFILL_LANES_PER_QUERY": [4, 8],
+    "CLUSTER_TILE_K_F32": [8, 16, 32],
+    "NO_SG_TILE_K_MAX": [8, 16]
+  },
+  "cases": [
+    {
+      "name": "gqa-f32-1x64x64-h4kv2d16",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 64, "kvSeq": 64, "heads": 4, "kvHeads": 2, "headDim": 16 },
+      "attrs": { "num_heads": 4, "kv_num_heads": 2, "scale": 0.25, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 64, 64], "dtype": "float32", "dist": "normal", "seed": 209, "scale": 0.2 },
+        "keyT": { "shape": [1, 64, 32], "dtype": "float32", "dist": "normal", "seed": 210, "scale": 0.2 },
+        "valueT": { "shape": [1, 64, 32], "dtype": "float32", "dist": "normal", "seed": 211, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 64, 64], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 2, 64, 16], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 2, 64, 16], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "qwen3-prefill-32h8kv-d128-s512",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 512, "kvSeq": 512, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 512, 4096], "dtype": "float32", "dist": "normal", "seed": 370, "scale": 0.2 },
+        "keyT": { "shape": [1, 512, 1024], "dtype": "float32", "dist": "normal", "seed": 371, "scale": 0.2 },
+        "valueT": { "shape": [1, 512, 1024], "dtype": "float32", "dist": "normal", "seed": 372, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 512, 4096], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 512, 128], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 512, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "decode-h4kv1-d64-kv512",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 1, "kvSeq": 512, "heads": 4, "kvHeads": 1, "headDim": 64 },
+      "attrs": { "num_heads": 4, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 376, "scale": 0.2 },
+        "keyT": { "shape": [1, 512, 64], "dtype": "float32", "dist": "normal", "seed": 377, "scale": 0.2 },
+        "valueT": { "shape": [1, 512, 64], "dtype": "float32", "dist": "normal", "seed": 378, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 256], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 512, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 512, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * args.kvSeq * args.kvHeads * args.headDim + 2 * args.qSeq * args.heads * args.headDim) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "llama-decode-32h8kv-d128-kv2048",
+      "preset": "model",
+      "vars": { "batch": 1, "qSeq": 1, "kvSeq": 2048, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 4096], "dtype": "float32", "dist": "normal", "seed": 373, "scale": 0.2 },
+        "keyT": { "shape": [1, 2048, 1024], "dtype": "float32", "dist": "normal", "seed": 374, "scale": 0.2 },
+        "valueT": { "shape": [1, 2048, 1024], "dtype": "float32", "dist": "normal", "seed": 375, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 4096], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 2048, 128], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 2048, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * args.kvSeq * args.kvHeads * args.headDim + 2 * args.qSeq * args.heads * args.headDim) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "llama-decode-32h8kv-d128-kv8192",
+      "preset": "stress",
+      "vars": { "batch": 1, "qSeq": 1, "kvSeq": 8192, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 4096], "dtype": "float32", "dist": "normal", "seed": 373, "scale": 0.2 },
+        "keyT": { "shape": [1, 8192, 1024], "dtype": "float32", "dist": "normal", "seed": 374, "scale": 0.2 },
+        "valueT": { "shape": [1, 8192, 1024], "dtype": "float32", "dist": "normal", "seed": 375, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [8191] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [8192] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 4096], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 8192, 128], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 8192, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": false,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * args.kvSeq * args.kvHeads * args.headDim + 2 * args.qSeq * args.heads * args.headDim) * 4"
+          }
+        ]
+      },
+      "provenance": {
+        "notes": "Long-context decode beyond the MAX_SPLITS boundary: ceilDiv(kvSeq, QKV_SPLIT_TILE_K)=64 exceeds the cap of 16. Paired with the headDim-32 control, this case checks the tradeoff between additional split parallelism and merge overhead without assuming that one flat cap is optimal for both geometries."
+      }
+    },
+    {
+      "name": "qwen3-prefill-32h8kv-d128-s512-f16",
+      "preset": "model",
+      "vars": { "batch": 1, "qSeq": 512, "kvSeq": 512, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 512, 4096], "dtype": "float16", "dist": "normal", "seed": 370, "scale": 0.2 },
+        "keyT": { "shape": [1, 512, 1024], "dtype": "float16", "dist": "normal", "seed": 371, "scale": 0.2 },
+        "valueT": { "shape": [1, 512, 1024], "dtype": "float16", "dist": "normal", "seed": 372, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 512, 4096], "dtype": "float16" },
+        "presentKeyT": { "shape": [1, 8, 512, 128], "dtype": "float16" },
+        "presentValueT": { "shape": [1, 8, 512, 128], "dtype": "float16" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "llama-decode-32h8kv-d128-kv2048-f16",
+      "preset": "model",
+      "vars": { "batch": 1, "qSeq": 1, "kvSeq": 2048, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 4096], "dtype": "float16", "dist": "normal", "seed": 373, "scale": 0.2 },
+        "keyT": { "shape": [1, 2048, 1024], "dtype": "float16", "dist": "normal", "seed": 374, "scale": 0.2 },
+        "valueT": { "shape": [1, 2048, 1024], "dtype": "float16", "dist": "normal", "seed": 375, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 4096], "dtype": "float16" },
+        "presentKeyT": { "shape": [1, 8, 2048, 128], "dtype": "float16" },
+        "presentValueT": { "shape": [1, 8, 2048, 128], "dtype": "float16" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * args.kvSeq * args.kvHeads * args.headDim + 2 * args.qSeq * args.heads * args.headDim) * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-decode-32h8kv-d128-kv2048",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 2048, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 2048, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 1) * dim(shapes.pastKeyT, 2) * (dim(shapes.queryT, 2) / attrs.num_heads) * attrs.num_heads"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-32h8kv-d128-s64",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 64, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 64, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 64, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [63] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 64, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 64, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 64, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-32h8kv-d128-s128",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 128, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 128, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [127] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 128, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 128, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-32h8kv-d128-s256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 256, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 256, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 256, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 256, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 256, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 256, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-32h8kv-d128-s512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-rotary-32h8kv-d128-s512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.013, "cosStep": 0.037 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-decode-f16-32h8kv-d128-kv2048",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 8, 2048, 128] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 8, 2048, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 1) * dim(shapes.pastKeyT, 2) * (dim(shapes.queryT, 2) / attrs.num_heads) * attrs.num_heads"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-f16-32h8kv-d128-s512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-prefill-rotary-f16-32h8kv-d128-s512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "cosCacheT": {
+          "dtype": "float16",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "sinCacheT": {
+          "dtype": "float16",
+          "shape": [512, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.013, "cosStep": 0.037 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float16", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "newkv-past-decode-q1p2048",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.013, "cosStep": 0.017 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.019, "cosStep": 0.023 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2049] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 2049, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 2049, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2) * (dim(shapes.queryT, 2) / attrs.num_heads) * attrs.num_heads"
+          }
+        ]
+      }
+    },
+    {
+      "name": "newkv-past-prefill-q512p512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.013, "cosStep": 0.017 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.019, "cosStep": 0.023 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.presentKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "quant-int8-decode-h8kv2-d128-kv2048",
+      "preset": "smoke",
+      "provenance": {
+        "notes": "Buffer-sharing INT8 decode. The bandwidth metric counts the portable physical minimum: WebGPU exposes each logical INT8 cache element through i32 storage, the past cache is preserved into present K/V, and attention reads K/V once per KV head. It intentionally does not credit duplicate cache reads across grouped query heads."
+      },
+      "attrs": {
+        "num_heads": 8,
+        "kv_num_heads": 2,
+        "kv_cache_bit_width": 8,
+        "k_quant_type": "PER_TENSOR",
+        "v_quant_type": "PER_TENSOR"
+      },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.013, "cosStep": 0.031 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.015, "cosStep": 0.033 }
+        },
+        "pastKeyT": { "dtype": "int8", "shape": [1, 2, 2048, 128], "dist": "uniform", "seed": 880, "min": -8, "max": 8 },
+        "pastValueT": {
+          "dtype": "int8",
+          "shape": [1, 2, 2048, 128],
+          "dist": "uniform",
+          "seed": 881,
+          "min": -8,
+          "max": 8
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } },
+        "kScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } },
+        "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 1024] },
+        "presentKeyT": { "dtype": "int8", "shape": [1, 2, 2048, 128] },
+        "presentValueT": { "dtype": "int8", "shape": [1, 2, 2048, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "4 * (6 * numel(shapes.pastKeyT) + 4 * numel(shapes.keyT) + 2 * numel(shapes.queryT))"
+          }
+        ]
+      }
+    },
+    {
+      "name": "quant-int8-prefill-h32kv8-d128-s512-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Production-sized prefill for the int8 KV-cache contract, retained as a performance witness for quantized-cache attention."
+      },
+      "attrs": {
+        "num_heads": 32,
+        "kv_num_heads": 8,
+        "kv_cache_bit_width": 8,
+        "k_quant_type": "PER_TENSOR",
+        "v_quant_type": "PER_TENSOR"
+      },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.013, "cosStep": 0.031 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.015, "cosStep": 0.033 }
+        },
+        "pastKeyT": { "dtype": "int8", "shape": [1, 8, 512, 128], "dist": "uniform", "seed": 882, "min": -8, "max": 8 },
+        "pastValueT": {
+          "dtype": "int8",
+          "shape": [1, 8, 512, 128],
+          "dist": "uniform",
+          "seed": 883,
+          "min": -8,
+          "max": 8
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "kScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } },
+        "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "int8", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "int8", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "quant-int4-prefill-h32kv8-d128-s512-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Production-sized prefill for the packed int4 KV-cache contract, retained as a performance witness for quantized-cache attention."
+      },
+      "attrs": {
+        "num_heads": 32,
+        "kv_num_heads": 8,
+        "kv_cache_bit_width": 4,
+        "k_quant_type": "PER_TENSOR",
+        "v_quant_type": "PER_TENSOR"
+      },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.013, "cosStep": 0.031 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.015, "cosStep": 0.033 }
+        },
+        "pastKeyT": { "dtype": "uint8", "shape": [1, 8, 512, 64], "dist": "uniform", "seed": 884, "min": 0, "max": 256 },
+        "pastValueT": {
+          "dtype": "uint8",
+          "shape": [1, 8, 512, 64],
+          "dist": "uniform",
+          "seed": 885,
+          "min": 0,
+          "max": 256
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "kScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } },
+        "vScaleT": { "dtype": "float32", "shape": [1], "data": { "kind": "values", "values": [0.0625] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "uint8", "shape": [1, 8, 512, 64] },
+        "presentValueT": { "dtype": "uint8", "shape": [1, 8, 512, 64] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "bias-prefill-cliff-h32kv8-d128-s512",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.05, "sinStep": 0.007, "cosStep": 0.009 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "hd16-belowfloor-decode-h8kv2-kv2048",
+      "preset": "smoke",
+      "attrs": { "num_heads": 8, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 32], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 32], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 2048, 16],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 2048, 16],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 128] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 2048, 16] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 2048, 16] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * dim(shapes.pastKeyT, 1) * dim(shapes.pastKeyT, 2) * dim(shapes.pastKeyT, 3) * 4 + 2 * dim(shapes.queryT, 2) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sliding-window-decode-h8kv2-d128-kv2048-w256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 8, "kv_num_heads": 2, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 256], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 256], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 2048, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2047] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2048] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 1024] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 2048, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 2048, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * dim(shapes.pastKeyT, 1) * dim(shapes.pastKeyT, 2) * dim(shapes.pastKeyT, 3) * 4 + 2 * dim(shapes.queryT, 2) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "scalar-fallback-decode-hd32-h8kv8-kv4096",
+      "preset": "stress",
+      "vars": { "batch": 1, "qSeq": 1, "kvSeq": 4096, "heads": 8, "kvHeads": 8, "headDim": 32 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8, "scale": 0.17677669529663687, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 911, "scale": 0.2 },
+        "keyT": { "shape": [1, 4096, 256], "dtype": "float32", "dist": "normal", "seed": 912, "scale": 0.2 },
+        "valueT": { "shape": [1, 4096, 256], "dtype": "float32", "dist": "normal", "seed": 913, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4095] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [4096] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "empty" },
+        "presentKeyT": { "shape": [1, 8, 4096, 32], "dtype": "float32", "dist": "empty" },
+        "presentValueT": { "shape": [1, 8, 4096, 32], "dtype": "float32", "dist": "empty" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "(4 * dim(shapes.keyT, 1) * attrs.kv_num_heads * (dim(shapes.queryT, 2) / attrs.num_heads) + 2 * dim(shapes.queryT, 2)) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "scalar-fallback-prefill-hd32-h8kv8-s512",
+      "preset": "stress",
+      "vars": { "batch": 1, "qSeq": 512, "kvSeq": 512, "heads": 8, "kvHeads": 8, "headDim": 32 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 8, "scale": 0.17677669529663687, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 512, 256], "dtype": "float32", "dist": "normal", "seed": 901, "scale": 0.2 },
+        "keyT": { "shape": [1, 512, 256], "dtype": "float32", "dist": "normal", "seed": 902, "scale": 0.2 },
+        "valueT": { "shape": [1, 512, 256], "dtype": "float32", "dist": "normal", "seed": 903, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 512, 256], "dtype": "float32", "dist": "empty" },
+        "presentKeyT": { "shape": [1, 8, 512, 32], "dtype": "float32", "dist": "empty" },
+        "presentValueT": { "shape": [1, 8, 512, 32], "dtype": "float32", "dist": "empty" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 0) * dim(shapes.queryT, 1) * dim(shapes.keyT, 1) * attrs.num_heads * (dim(shapes.queryT, 2) / attrs.num_heads)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "gqa-f32-q31-kv511-h8-kv2-d64-double-threshold-cliff",
+      "preset": "stress",
+      "vars": { "batch": 1, "qSeq": 31, "kvSeq": 511, "heads": 8, "kvHeads": 2, "headDim": 64 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 31, 512], "dtype": "float32", "dist": "normal", "seed": 921, "scale": 0.2 },
+        "keyT": { "shape": [1, 511, 128], "dtype": "float32", "dist": "normal", "seed": 922, "scale": 0.2 },
+        "valueT": { "shape": [1, 511, 128], "dtype": "float32", "dist": "normal", "seed": 923, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [510] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 31, 512], "dtype": "float32", "dist": "empty" },
+        "presentKeyT": { "shape": [1, 2, 511, 64], "dtype": "float32", "dist": "empty" },
+        "presentValueT": { "shape": [1, 2, 511, 64], "dtype": "float32", "dist": "empty" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "gqa-f32-q32-kv512-h8-kv2-d64-threshold-control",
+      "preset": "stress",
+      "vars": { "batch": 1, "qSeq": 32, "kvSeq": 512, "heads": 8, "kvHeads": 2, "headDim": 64 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 2, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 32, 512], "dtype": "float32", "dist": "normal", "seed": 924, "scale": 0.2 },
+        "keyT": { "shape": [1, 512, 128], "dtype": "float32", "dist": "normal", "seed": 925, "scale": 0.2 },
+        "valueT": { "shape": [1, 512, 128], "dtype": "float32", "dist": "normal", "seed": 926, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 32, 512], "dtype": "float32", "dist": "empty" },
+        "presentKeyT": { "shape": [1, 2, 512, 64], "dtype": "float32", "dist": "empty" },
+        "presentValueT": { "shape": [1, 2, 512, 64], "dtype": "float32", "dist": "empty" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "headsink-prefill-h32kv8-d128-s512-generic-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Llama-sized GQA prefill with smooth-softmax head sinks, measuring the additional score transformation on a production attention shape."
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "headSinkT": { "dtype": "float32", "shape": [32], "dist": "normal", "seed": 927, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "bias-headsink-prefill-h32kv8-d128-s512-generic-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "The valid additive-bias plus smooth-softmax head-sink combination has a distinct manifest variant but no flash-prefill specialization."
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 32, 512, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.05, "sinStep": 0.007, "cosStep": 0.009 }
+        },
+        "headSinkT": { "dtype": "float32", "shape": [32], "dist": "normal", "seed": 928, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "softcap-prefill-h32kv8-d128-s512-generic-pathology",
+      "preset": "stress",
+      "provenance": {
+        "notes": "Production prefill with model-used score soft-capping, measuring the capped-score semantic on the past-KV contract."
+      },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "softcap": 30 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 1024], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 512, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 512, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 512, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "4 * dim(shapes.queryT, 2) * (dim(shapes.queryT, 1) * (dim(shapes.pastKeyT, 2) - dim(shapes.queryT, 1)) + dim(shapes.queryT, 1) * (dim(shapes.queryT, 1) + 1) / 2)"
+          }
+        ]
+      }
+    },
+    {
+      "name": "qkv-prefill-f32-h4kv1-d64-s128",
+      "preset": "model",
+      "provenance": {
+        "notes": "Lower f32 register-topology boundary comparing the Q32-broadcast and subgroup-cluster routes."
+      },
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 128, "heads": 4, "kvHeads": 1, "headDim": 64 },
+      "attrs": { "num_heads": 4, "kv_num_heads": 1, "scale": 0.125, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 128, 256], "dtype": "float32", "dist": "normal", "seed": 1075, "scale": 0.2 },
+        "keyT": { "shape": [1, 128, 64], "dtype": "float32", "dist": "normal", "seed": 1076, "scale": 0.2 },
+        "valueT": { "shape": [1, 128, 64], "dtype": "float32", "dist": "normal", "seed": 1077, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [127] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 128, 256], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 128, 64], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 128, 64], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "qkv-prefill-f32-h4kv1-d256-s128",
+      "preset": "model",
+      "provenance": {
+        "notes": "Upper f32 register-topology boundary comparing Q32-broadcast, subgroup-cluster, and generic online routes."
+      },
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 128, "heads": 4, "kvHeads": 1, "headDim": 256 },
+      "attrs": { "num_heads": 4, "kv_num_heads": 1, "scale": 0.0625, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 128, 1024], "dtype": "float32", "dist": "normal", "seed": 1081, "scale": 0.2 },
+        "keyT": { "shape": [1, 128, 256], "dtype": "float32", "dist": "normal", "seed": 1082, "scale": 0.2 },
+        "valueT": { "shape": [1, 128, 256], "dtype": "float32", "dist": "normal", "seed": 1083, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [127] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 128, 1024], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 1, 128, 256], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 1, 128, 256], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "window-decode-32h8kv-d128-cap1024-w1024",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sliding_window_cache": 1, "local_window_size": 1024 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "window-decode-32h8kv-d128-cap1024-unfilled",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sliding_window_cache": 1, "local_window_size": 1024 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "window-decode-32h8kv-d128-cap256-w256",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 256, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 256, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 256, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 256, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "window-chunk-prefill-32h8kv-d128-cap1024-q128",
+      "preset": "smoke",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sliding_window_cache": 1, "local_window_size": 1024 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" },
+          {
+            "type": "bandwidth",
+            "name": "cache traffic",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      },
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 1024, "heads": 32, "headDim": 128 }
+    },
+    {
+      "name": "window-decode-32h8kv-d128-cap4096-w4096",
+      "preset": "model",
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "sliding_window_cache": 1, "local_window_size": 4096 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 1, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 4096, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 4096, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65535] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [65536] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 1, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 4096, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 4096, 128] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "2 * 2 * numel(shapes.pastKeyT) * 4 + numel(shapes.queryT) * 4 + numel(shapes.outputT) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-chunk-append-32h8kv-d128-cap1024-q128",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 1024, "heads": 32, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 4096],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 1024],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 8, 1024, 128],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 4096] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 8, 1024, 128] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 8, 1024, 128] }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
+    },
+    {
+      "name": "sharedkv-chunk-append-8h2kv-d256-cap1024-q128",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 128, "kvSeq": 1024, "heads": 8, "headDim": 256 },
+      "attrs": { "num_heads": 8, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 2048],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.011, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.031, "cosStep": 0.007 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 512],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.037, "cosStep": 0.005 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1024, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.017, "cosStep": 0.019 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 1024, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.023, "cosStep": 0.013 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 2048] },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 1024, 256] },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 1024, 256] }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      },
+      "provenance": {
+        "source": "register-geometry gate asymmetry",
+        "notes": "Cached f32 prefill at headDim 256, above the shared-memory cluster's register-geometry boundary. This guards consistent route admission between no-past and cached prefill families."
+      }
+    }
+  ]
+}

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

	@@ -0,0 +1,285 @@

+{% set usesF16 = usesF16 | default(false) %}
+{% set quantized = quantized | default(false) %}
+{% set bits = bits | default(0) %}
+{% if useSeqlens is not defined %}{% set useSeqlens = false %}{% endif %}
+{% if usesF16 %}enable f16;
+{% endif %}{{ env.wgsl.resourceDeclarations }}
+// f16 inputs widen to f32 on load; computation and shared memory remain f32,
+// and results narrow only on store. The casts are identities for f32.
+{% set IO = "f16" if usesF16 else "f32" %}
+{% if cooperative %}
+// One workgroup owns each (batch, qHead, qToken). Lanes partition the key range,
+// run local online softmax states, then merge max/rescale/sum through workgroup
+// memory. This exposes parallelism when the query count is small.
+{% else %}
+// One invocation owns each (batch, qHead, qToken) and walks the key range. The
+// larger query grid supplies parallelism without a workgroup-level merge.
+{% endif %}
+// The attention path supports optional per-head qk-norm and
+// rotary on Q, causal (+ sliding-window) masking, logit soft-capping, additive
+// attention bias, and smooth-softmax / head-sink. The cache is read directly
+// (f32), or dequantized from standard INT8 / +8-biased packed INT4 storage when
+// `quantized` is enabled.
+const HEAD_DIM: u32 = {{ headDim }}u;
+{% if hasRotary %}const HALF: u32 = {{ half }}u;{% endif %}
+const Q_HEADS: u32 = {{ qHeads }}u;
+const KV_HEADS: u32 = {{ kvHeads }}u;
+const Q_HIDDEN: u32 = {{ qHidden }}u;
+const GROUP: u32 = {{ qHeads }}u / {{ kvHeads }}u;
+const PACKED: u32 = {{ packed }}u;
+{% if hasQNorm %}const QK_EPS: f32 = {{ qkEps }};{% endif %}
+{% if cooperative %}const WG: u32 = {{ tunables.COOPERATIVE_WORKGROUP_SIZE }}u;{% else %}const WG: u32 = {{ scalarWorkgroupSize }}u;{% endif %}
+const NEG_INF: f32 = -3.4028234663852886e38;
+{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
+fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+{% if quantized %}
+fn kscale(d: u32, hk: u32) -> f32 { return k_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
+fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
+{% endif %}
+fn read_k(base: u32, d: u32{% if quantized %}, hk: u32{% endif %}) -> f32 {
+{% if not quantized %}
+  return f32(kcache[base + d]);
+{% elif bits == 8 %}
+  return f32(kcache[base + d]) * kscale(d, hk);
+{% else %}
+  let byte = kcache[base + (d >> 1u)];
+  let nib = i32((byte >> ((d & 1u) * 4u)) & 0xFu) - 8;
+  return f32(nib) * kscale(d, hk);
+{% endif %}
+}
+fn read_v(base: u32, d: u32{% if quantized %}, hk: u32{% endif %}) -> f32 {
+{% if not quantized %}
+  return f32(vcache[base + d]);
+{% elif bits == 8 %}
+  return f32(vcache[base + d]) * vscale(d, hk);
+{% else %}
+  let byte = vcache[base + (d >> 1u)];
+  let nib = i32((byte >> ((d & 1u) * 4u)) & 0xFu) - 8;
+  return f32(nib) * vscale(d, hk);
+{% endif %}
+}
+{% if cooperative %}
+var<workgroup> qsh: array<f32, HEAD_DIM>;       // query head-vector (post norm/rotary)
+var<workgroup> acc_sh: array<f32, HEAD_DIM * 32u>; // per-lane V accumulator [lane*HEAD_DIM + d]
+var<workgroup> m_sh: array<f32, 32u>;
+var<workgroup> l_sh: array<f32, 32u>;
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  // One workgroup per query; fold large query grids over x/y.
+  let qi = wid.x + wid.y * nwg.x;
+  let total = params.batch * Q_HEADS * params.qSeq;
+  if (qi >= total) { return; }
+  let lane = lid.x;
+  let s = qi % params.qSeq;
+  let tmp = qi / params.qSeq;
+  let h = tmp % Q_HEADS;
+  let b = tmp / Q_HEADS;
+  let hk = h / GROUP;
+  let totalSeq = params.totalSeq;
+{% if useSeqlens %}
+  // Buffer-sharing caches keep capacity in the BNSH dim2 stride; seqlens_k[b]+1
+  // is the per-batch active end and rows beyond it are untouched storage.
+  let activeEnd = min(totalSeq, max(params.qSeq, u32(seqlens_k[b]) + 1u));
+{% else %}
+  let activeEnd = totalSeq;
+{% endif %}
+  let absPos = activeEnd - params.qSeq + s;
+  // Cooperative query-vector prep (post norm/rotary) into shared memory.
+  let qBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { qsh[d] = f32(query[qBase + d]); }
+  workgroupBarrier();
+{% if hasQNorm %}
+  var part = 0.0;
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { part = part + qsh[d] * qsh[d]; }
+  m_sh[lane] = part;
+  workgroupBarrier();
+  var ms = 0.0;
+  for (var L = 0u; L < WG; L = L + 1u) { ms = ms + m_sh[L]; }
+  let invRms = inverseSqrt(ms / f32(HEAD_DIM) + QK_EPS);
+  for (var d = lane; d < HEAD_DIM; d = d + WG) { qsh[d] = qsh[d] * invRms * f32(q_norm_weight[d]); }
+  workgroupBarrier();
+{% endif %}
+{% if hasRotary %}
+  // Each lane owns disjoint pairs (d, d+HALF), so the in-place rotate is safe.
+  for (var d = lane; d < HALF; d = d + WG) {
+    let cs = f32(cos_cache[absPos * HALF + d]);
+    let sn = f32(sin_cache[absPos * HALF + d]);
+    let x0 = qsh[d];
+    let x1 = qsh[d + HALF];
+    qsh[d] = x0 * cs - x1 * sn;
+    qsh[d + HALF] = x1 * cs + x0 * sn;
+  }
+  workgroupBarrier();
+{% endif %}
+  let scale = scale_value();
+  let maxKj = absPos + 1u;
+  var minKj = 0u;
+  if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
+  let accBase = lane * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = 0.0; }
+  var m = NEG_INF;
+  var l = 0.0;
+  // Each lane sweeps a strided slice of the key range.
+  for (var j = minKj + lane; j < maxKj; j = j + WG) {
+    let base = ((b * KV_HEADS + hk) * totalSeq + j) * PACKED;
+    var dot = 0.0;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { dot = dot + qsh[d] * read_k(base, d{% if quantized %}, hk{% endif %}); }
+    var score = dot * scale;
+    if (params.softcap > 0.0) { score = params.softcap * tanh(clamp(score / params.softcap, -30.0, 30.0)); }
+{% if hasBias %}
+    let bb = select(b, 0u, params.biasBatch == 1u);
+    let bh = select(h, 0u, params.biasHeads == 1u);
+    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * totalSeq + j]);
+{% endif %}
+    let newM = max(m, score);
+    let corr = exp(m - newM);
+    let p = exp(score - newM);
+    l = l * corr + p;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = acc_sh[accBase + d] * corr + p * read_v(base, d{% if quantized %}, hk{% endif %}); }
+    m = newM;
+  }
+  // Flash merge across the WG lanes: global max, rescale, summed denom.
+  m_sh[lane] = m;
+  workgroupBarrier();
+  var gm = NEG_INF;
+  for (var L = 0u; L < WG; L = L + 1u) { gm = max(gm, m_sh[L]); }
+  let fctr = exp(m - gm);
+  l = l * fctr;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc_sh[accBase + d] = acc_sh[accBase + d] * fctr; }
+  l_sh[lane] = l;
+  workgroupBarrier();
+  var glsum = 0.0;
+  for (var L = 0u; L < WG; L = L + 1u) { glsum = glsum + l_sh[L]; }
+  var sink = 0.0;
+{% if hasHeadSink %}
+  sink = f32(head_sink[h]);
+  let useSmooth = true;
+{% else %}
+  let useSmooth = params.smoothSoftmax != 0u;
+{% endif %}
+  var finalM = gm;
+  var smoothTerm = 0.0;
+  if (useSmooth) {
+    finalM = max(gm, sink);
+    smoothTerm = exp(sink - finalM);
+  }
+  let accScale = exp(gm - finalM);
+  let invDenom = accScale / (smoothTerm + glsum * accScale);
+  // Cross-lane sum of the rescaled V accumulators, written cooperatively.
+  let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  for (var d = lane; d < HEAD_DIM; d = d + WG) {
+    var sumacc = 0.0;
+    for (var L = 0u; L < WG; L = L + 1u) { sumacc = sumacc + acc_sh[L * HEAD_DIM + d]; }
+    output[oBase + d] = {{ IO }}(sumacc * invDenom);
+  }
+}
+{% else %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  // One thread per query; the grid is 2D-folded so the workgroup count stays
+  // within maxComputeWorkgroupsPerDimension for large prefills. Reconstruct the
+  // flat thread index from the folded workgroup grid (wid.x + wid.y*nwg.x) — the
+  // global_invocation_id only spans nwg.x*WG and would drop the y-folded tail.
+  // This reduces to the ordinary flat index when nwg.y == 1.
+  let qi = (wid.x + wid.y * nwg.x) * WG + lid.x;
+  let total = params.batch * Q_HEADS * params.qSeq;
+  if (qi >= total) { return; }
+  let s = qi % params.qSeq;
+  let tmp = qi / params.qSeq;
+  let h = tmp % Q_HEADS;
+  let b = tmp / Q_HEADS;
+  let hk = h / GROUP;
+  let totalSeq = params.totalSeq;
+{% if useSeqlens %}
+  // Buffer-sharing caches keep capacity in the BNSH dim2 stride; seqlens_k[b]+1
+  // is the per-batch active end and rows beyond it are untouched storage.
+  let activeEnd = min(totalSeq, max(params.qSeq, u32(seqlens_k[b]) + 1u));
+{% else %}
+  let activeEnd = totalSeq;
+{% endif %}
+  let absPos = activeEnd - params.qSeq + s;
+  // Private query-vector prep (post norm/rotary).
+  var q: array<f32, HEAD_DIM>;
+  let qBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = f32(query[qBase + d]); }
+{% if hasQNorm %}
+  var ss = 0.0;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ss = ss + q[d] * q[d]; }
+  let invRms = inverseSqrt(ss / f32(HEAD_DIM) + QK_EPS);
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = q[d] * invRms * f32(q_norm_weight[d]); }
+{% endif %}
+{% if hasRotary %}
+  for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = f32(cos_cache[absPos * HALF + d]);
+    let sn = f32(sin_cache[absPos * HALF + d]);
+    let x0 = q[d];
+    let x1 = q[d + HALF];
+    q[d] = x0 * cs - x1 * sn;
+    q[d + HALF] = x1 * cs + x0 * sn;
+  }
+{% endif %}
+  let scale = scale_value();
+  let maxKj = absPos + 1u;
+  var minKj = 0u;
+  if (params.windowSize > 0u && maxKj > params.windowSize) { minKj = maxKj - params.windowSize; }
+  var acc: array<f32, HEAD_DIM>;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc[d] = 0.0; }
+  var m = NEG_INF;
+  var l = 0.0;
+  for (var j = minKj; j < maxKj; j = j + 1u) {
+    let base = ((b * KV_HEADS + hk) * totalSeq + j) * PACKED;
+    var dot = 0.0;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { dot = dot + q[d] * read_k(base, d{% if quantized %}, hk{% endif %}); }
+    var score = dot * scale;
+    if (params.softcap > 0.0) { score = params.softcap * tanh(clamp(score / params.softcap, -30.0, 30.0)); }
+{% if hasBias %}
+    let bb = select(b, 0u, params.biasBatch == 1u);
+    let bh = select(h, 0u, params.biasHeads == 1u);
+    score = score + f32(attn_bias[((bb * params.biasHeads + bh) * params.qSeq + s) * totalSeq + j]);
+{% endif %}
+    let newM = max(m, score);
+    let corr = exp(m - newM);
+    let p = exp(score - newM);
+    l = l * corr + p;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) { acc[d] = acc[d] * corr + p * read_v(base, d{% if quantized %}, hk{% endif %}); }
+    m = newM;
+  }
+  var sink = 0.0;
+{% if hasHeadSink %}
+  sink = f32(head_sink[h]);
+  let useSmooth = true;
+{% else %}
+  let useSmooth = params.smoothSoftmax != 0u;
+{% endif %}
+  var finalM = m;
+  var smoothTerm = 0.0;
+  if (useSmooth) {
+    finalM = max(m, sink);
+    smoothTerm = exp(sink - finalM);
+  }
+  let accScale = exp(m - finalM);
+  let invDenom = accScale / (smoothTerm + l * accScale);
+  let oBase = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { output[oBase + d] = {{ IO }}(acc[d] * invDenom); }
+}
+{% endif %}

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

	@@ -0,0 +1,440 @@

+{% if source is defined and source.mode is defined %}
+{% set mode = source.mode %}
+{% elif mode is not defined %}
+{% set mode = "transpose" %}
+{% endif %}
+{% set presentScalar = presentScalar | default(inputScalar | default("f32")) %}
+{% set presentElement = presentElement | default(presentScalar) %}
+{% set packed = packed | default(0) %}
+{% set bits = bits | default(0) %}
+{% set qmax = qmax | default(0) %}
+{% set qmin = qmin | default(0) %}
+{% if usesF16 is defined and usesF16 %}enable f16;
+{% endif %}{{ env.wgsl.resourceDeclarations }}
+{% if mode == "transpose" %}
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+const WG: u32 = {{ copyWorkgroupSize }}u;
+{% if presentVec4 %}
+// Transpose copy [batch, token, head*headDim] -> present [batch, head, token,
+// headDim]. The headDim row is contiguous on both sides and divisible by four,
+// so each invocation copies a vec4 along d with coalesced stores.
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
+{% else %}
+const HEAD_DIM: u32 = {{ headDim }}u;
+const KV_HIDDEN: u32 = {{ kvHidden }}u;
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
+  // flat invocation index; this reduces to gid.x when no fold is needed.
+  let index = gid.x + gid.y * nwg.x * WG;
+{% if presentVec4 %}
+  let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM_V4;
+  if (index >= total) {
+    return;
+  }
+  let d4 = index % HEAD_DIM_V4;
+  let token = (index / HEAD_DIM_V4) % params.kvSeq;
+  let head = (index / (HEAD_DIM_V4 * params.kvSeq)) % KV_HEADS;
+  let batch = index / (HEAD_DIM_V4 * params.kvSeq * KV_HEADS);
+  let packed = (batch * params.kvSeq + token) * KV_HIDDEN_V4 + head * HEAD_DIM_V4 + d4;
+  present_key[index] = {{ presentElement }}(key[packed]);
+  present_value[index] = {{ presentElement }}(value[packed]);
+{% else %}
+  let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM;
+  if (index >= total) {
+    return;
+  }
+  let d = index % HEAD_DIM;
+  let token = (index / HEAD_DIM) % params.kvSeq;
+  let head = (index / (HEAD_DIM * params.kvSeq)) % KV_HEADS;
+  let batch = index / (HEAD_DIM * params.kvSeq * KV_HEADS);
+  let packed_index = (batch * params.kvSeq + token) * KV_HIDDEN + head * HEAD_DIM + d;
+  present_key[index] = {{ presentScalar }}(key[packed_index]);
+  present_value[index] = {{ presentScalar }}(value[packed_index]);
+{% endif %}
+}
+{%- else %}
+{% set shareAppend = mode == "merge_share" and shareRegion == "append" %}
+{% set cooperativeCopy = device.adapterInfo.vendor != "apple" %}
+{% set cooperativeMerge = cooperativeCopy or shareAppend %}
+{% set cooperativeMode = cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
+{% set needsGid = not cooperativeMode or (shareAppend and not cooperativeCopy) %}
+{% if mode == "copy" %}
+// Copy the existing BNSH cache unchanged. Compression changes only the number
+// of stored words per row.
+{% elif mode == "merge" %}
+// Copy the past BNSH cache, then append new BSH K/V tokens into the present cache.
+{% elif mode == "merge_share" %}
+// Buffer-sharing append: past and present share the full-capacity BNSH stride.
+// Rows outside each batch's append window [seqlens_k[b]+1-keySeq, seqlens_k[b]+1)
+// are an identity copy of the past cache; rows inside it come from the new BSH
+// K/V (K optionally rotated at its absolute position).
+// The append and retain passes own disjoint ranges, so their order is immaterial.
+{% elif mode == "build" %}
+// Build a BNSH cache from new BSH K/V, optionally applying K RMSNorm and rotary.
+{% elif mode == "window_shift" %}
+// Windowed cache: compact the surviving past rows down by the eviction count and
+// append the new BSH K/V, keeping the most recent min(T, capacity) tokens
+// contiguous at rows [0, L). Reads past, writes present (distinct buffers).
+{% elif mode == "append_quant" %}
+// Quantize new BSH K/V into an existing buffer-sharing INT8 BNSH cache. The
+// active sequence end comes from seqlens_k; the physical cache stride remains
+// the full present-cache capacity.
+{% else %}
+// Quantize new BSH K/V into an INT8 or packed INT4 BNSH cache.
+{% endif %}
+const HEAD_DIM: u32 = {{ headDim }}u;
+{% if mode != "copy" %}const KV_HEADS: u32 = {{ kvHeads }}u;
+{% endif %}
+const WG: u32 = {{ tunables.COPY_WORKGROUP_SIZE }}u;
+{% if mode != "copy" and (mode != "merge_share" or shareAppend) %}const KV_HIDDEN: u32 = {{ kvHeads }}u * {{ headDim }}u;
+{% endif %}
+{% if mode == "build_quant" or mode == "append_quant" %}const PACKED: u32 = {{ packed }}u;
+{% endif %}
+{% if hasRotary and mode != "copy" and (mode != "merge_share" or shareAppend) %}const HALF: u32 = {{ half }}u;
+{% endif %}
+{% if hasKNorm and mode != "copy" and (mode != "merge_share" or shareAppend) %}const QK_EPS: f32 = {{ qkEps }};
+{% endif %}
+{% if mode == "build_quant" or mode == "append_quant" %}
+const QMAX: f32 = {{ qmax }};
+const QMIN: f32 = {{ qmin }};
+fn kscale(d: u32, hk: u32) -> f32 { return k_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
+fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
+{% endif %}
+@compute @workgroup_size(WG)
+fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>,
+        {% endif %}{% if cooperativeMode %}@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>,
+        {% endif %}
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+{% macro append_cooperative_walk() %}
+  // Compact cooperative append walk: the thread range spans only the appended
+  // window (batch x kvHead x keySeq rows), not the full cache capacity; the
+  // destination row is the dynamic append offset plus the local row.
+  // The loop is bounded by the elements this workgroup owns rather than a fixed
+  // HEAD_DIM trip count.
+  let wgEnd = min(elemBase + WG * HEAD_DIM, totalElems);
+  for (var e = elemBase + lid.x; e < wgEnd; e = e + WG) {
+    {
+      let d = e % HEAD_DIM;
+      let et = e / HEAD_DIM;
+      let j = et % params.keySeq;
+      let etmp = et / params.keySeq;
+      let hk = etmp % KV_HEADS;
+      let b = etmp / KV_HEADS;
+      // Clamping below by keySeq maps a right-padded first prompt
+      // (seqlens_k[b]+1 < keySeq) to append offset 0.
+      let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
+      let t = activeEnd - params.keySeq + j;
+      let dst = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM + d;
+      let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
+{% if hasRotary %}
+      // Element-wise NeoX rotation: the pair (dr, dr+HALF) shares one
+      // cos/sin entry, and row t is its own absolute position.
+      let dr = select(d - HALF, d, d < HALF);
+      let cs = cos_cache[t * HALF + dr];
+      let sn = sin_cache[t * HALF + dr];
+      let x0 = new_k[nSrc + dr];
+      let x1 = new_k[nSrc + dr + HALF];
+      present_key[dst] = select(x1 * cs + x0 * sn, x0 * cs - x1 * sn, d < HALF);
+{% else %}
+      present_key[dst] = new_k[nSrc + d];
+{% endif %}
+      present_value[dst] = new_v[nSrc + d];
+    }
+  }
+{% endmacro %}
+{% macro append_per_row_walk() %}
+  // Per-row append walk: one thread streams its row's contiguous bytes.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let j = i % params.keySeq;
+  let tmp = i / params.keySeq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  // Clamping below by keySeq maps a right-padded first prompt
+  // (seqlens_k[b]+1 < keySeq) to append offset 0.
+  let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
+  let t = activeEnd - params.keySeq + j;
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
+{% if hasRotary %}
+  var k: array<f32, HEAD_DIM>;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = f32(new_k[nSrc + d]); }
+  // Row t is its own absolute position in the shared cache.
+  for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = cos_cache[t * HALF + d];
+    let sn = sin_cache[t * HALF + d];
+    let x0 = k[d];
+    let x1 = k[d + HALF];
+    k[d] = x0 * cs - x1 * sn;
+    k[d + HALF] = x1 * cs + x0 * sn;
+  }
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[dstBase + d] = k[d];
+    present_value[dstBase + d] = new_v[nSrc + d];
+  }
+{% else %}
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[dstBase + d] = new_k[nSrc + d];
+    present_value[dstBase + d] = new_v[nSrc + d];
+  }
+{% endif %}
+{% endmacro %}
+{% if cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
+  // Each workgroup copies the contiguous element span of its WG rows
+  // cooperatively: at step s lane l touches element s*WG + l, so adjacent
+  // lanes hit adjacent addresses on both source and destination. Same
+  // dispatch geometry (ceil(rows/WG) workgroups, folded into x/y).
+  let wgFlat = wid.x + wid.y * nwg.x;
+  let elemBase = wgFlat * (WG * HEAD_DIM);
+  let totalElems = params.count * HEAD_DIM;
+{% endif %}
+{% if mode == "copy" and cooperativeCopy %}
+  // Copy mode is an identity copy over the whole BNSH range: dst index == src index.
+  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
+    let e = elemBase + s * WG + lid.x;
+    if (e < totalElems) {
+      present_key[e] = src_k[e];
+      present_value[e] = src_v[e];
+    }
+  }
+{% elif mode == "merge" and cooperativeCopy %}
+  // present[0:pastSeq] = past (BNSH copy); present[pastSeq:totalSeq] = new K/V (BSH appended).
+  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
+    let e = elemBase + s * WG + lid.x;
+    if (e < totalElems) {
+      let d = e % HEAD_DIM;
+      let et = e / HEAD_DIM;
+      let t = et % params.seq;
+      let etmp = et / params.seq;
+      let hk = etmp % KV_HEADS;
+      let b = etmp / KV_HEADS;
+      if (t < params.pastSeq) {
+        let src = ((b * KV_HEADS + hk) * params.pastSeq + t) * HEAD_DIM + d;
+        present_key[e] = past_k[src];
+        present_value[e] = past_v[src];
+      } else {
+        let nSrc = (b * params.keySeq + (t - params.pastSeq)) * KV_HIDDEN + hk * HEAD_DIM + d;
+        present_key[e] = new_k[nSrc];
+        present_value[e] = new_v[nSrc];
+      }
+    }
+  }
+{% elif mode == "merge_share" and cooperativeMerge %}
+{% if shareAppend %}
+{% if cooperativeCopy %}
+{{ append_cooperative_walk() }}
+{% else %}
+  if (params.count < WG) {
+{{ append_cooperative_walk() }}
+    return;
+  }
+{{ append_per_row_walk() }}
+{% endif %}
+{% else %}
+  // Cooperative element walk (see the copy branch); the same-stride layout
+  // makes the outside-window rows an identity copy (dst index == src index).
+  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
+    let e = elemBase + s * WG + lid.x;
+    if (e < totalElems) {
+      let et = e / HEAD_DIM;
+      let t = et % params.seq;
+      let b = et / (params.seq * KV_HEADS);
+      // Clamping below by keySeq maps a right-padded first prompt
+      // (seqlens_k[b]+1 < keySeq) to append offset 0.
+      let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
+      let appendStart = activeEnd - params.keySeq;
+      // Everything outside the append window. The append pass owns the window itself,
+      // and the shared full-capacity stride makes this an identity copy (dst == src).
+      if (t < appendStart || t >= activeEnd) {
+        present_key[e] = past_k[e];
+        present_value[e] = past_v[e];
+      }
+    }
+  }
+{% endif %}
+{% elif mode == "copy" %}
+  // Thread i streams one contiguous (batch, kvHead, token) row.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let base = i * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[base + d] = src_k[base + d];
+    present_value[base + d] = src_v[base + d];
+  }
+{% elif mode == "merge" %}
+  // Per-row merge walk.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let t = i % params.seq;
+  let tmp = i / params.seq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  // present[0:pastSeq] = past (BNSH copy); present[pastSeq:totalSeq] = new K/V (BSH appended).
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  if (t < params.pastSeq) {
+    let srcBase = ((b * KV_HEADS + hk) * params.pastSeq + t) * HEAD_DIM;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = past_k[srcBase + d];
+      present_value[dstBase + d] = past_v[srcBase + d];
+    }
+  } else {
+    let nSrc = (b * params.keySeq + (t - params.pastSeq)) * KV_HIDDEN + hk * HEAD_DIM;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = new_k[nSrc + d];
+      present_value[dstBase + d] = new_v[nSrc + d];
+    }
+  }
+{% elif mode == "merge_share" %}
+  // Past and present share the full-capacity stride, so outside-window rows
+  // copy at the same index.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let t = i % params.seq;
+  let tmp = i / params.seq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  // Clamping below by keySeq maps a right-padded first prompt
+  // (seqlens_k[b]+1 < keySeq) to append offset 0.
+  let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
+  let appendStart = activeEnd - params.keySeq;
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  // Everything outside the append window. The append pass owns the window itself,
+  // and the shared full-capacity stride makes this an identity copy (dst == src).
+  if (t < appendStart || t >= activeEnd) {
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = past_k[dstBase + d];
+      present_value[dstBase + d] = past_v[dstBase + d];
+  }
+  }
+{% elif mode == "window_shift" %}
+  // Windowed cache: `params.seq` is a fixed capacity C, and the invariant is that
+  // the L = min(T, C) most recent tokens live contiguously at rows [0, L).
+  //
+  //   residentBefore = min(T - S, C)      E = max(0, residentBefore + S - C)
+  //   rows [0, appendStart)          <- past rows shifted down by E
+  //   rows [appendStart, append+S)   <- the new tokens
+  //   rows beyond that               <- outside the window, cleared
+  //
+  // The shift reads past and writes present, which must be distinct buffers: an
+  // in-place compaction would have one invocation overwrite row t while another
+  // still needs it as the source for row t-E, with no ordering between them.
+  //
+  // Attention needs no change for this layout. It derives its key range from
+  // `min(capacity, seqlens_k[b] + 1)`, which is exactly L, and both the causal
+  // and local-window masks depend only on the query/key distance:
+  //   q_abs - k_abs = (T - qSeq + s) - (origin + t) = (L - qSeq + s) - t
+  // so scoring a windowed cache as if it were a full L-length one is the same
+  // arithmetic. Only RoPE would need the true absolute position, which is why
+  // `windowShiftOk` refuses a rotary request outright rather than silently
+  // rotating at the cache row.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let t = i % params.seq;
+  let tmp = i / params.seq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  let absTotal = max(params.keySeq, u32(seqlens_k[b]) + 1u);
+  let residentBefore = min(absTotal - params.keySeq, params.seq);
+  // `max(0u, a + b - c)` does not clamp in u32: the subtraction wraps first, so a
+  // step that evicts nothing reads back ~2^32 instead of 0. Compare before
+  // subtracting.
+  let filled = residentBefore + params.keySeq;
+  let evicted = select(0u, filled - params.seq, filled > params.seq);
+  let appendStart = residentBefore - evicted;
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  if (t < appendStart) {
+    // Surviving past rows slide down by the eviction count.
+    let srcBase = ((b * KV_HEADS + hk) * params.seq + (t + evicted)) * HEAD_DIM;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = past_k[srcBase + d];
+      present_value[dstBase + d] = past_v[srcBase + d];
+    }
+  } else if (t < appendStart + params.keySeq) {
+    let nSrc = (b * params.keySeq + (t - appendStart)) * KV_HIDDEN + hk * HEAD_DIM;
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = new_k[nSrc + d];
+      present_value[dstBase + d] = new_v[nSrc + d];
+    }
+  } else {
+    // Outside the resident window. Cleared rather than left stale so the present
+    // buffer does not expose stale contents from its distinct allocation.
+    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+      present_key[dstBase + d] = {{ zeroScalar }}(0.0);
+      present_value[dstBase + d] = {{ zeroScalar }}(0.0);
+    }
+  }
+{% else %}
+  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
+  // flat invocation index; this reduces to gid.x when no fold is needed.
+  let i = gid.x + gid.y * nwg.x * WG;
+  if (i >= params.count) { return; }
+  let t = i % params.seq;
+  let tmp = i / params.seq;
+  let hk = tmp % KV_HEADS;
+  let b = tmp / KV_HEADS;
+  let srcBase = (b * params.seq + t) * KV_HIDDEN + hk * HEAD_DIM; // BSH new K/V
+  var k: array<f32, HEAD_DIM>;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = src_k[srcBase + d]; }
+{% if hasKNorm %}
+  var ms = 0.0;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ms = ms + k[d] * k[d]; }
+  let invRms = inverseSqrt(ms / f32(HEAD_DIM) + QK_EPS);
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = k[d] * invRms * k_norm_weight[d]; }
+{% endif %}
+{% if hasRotary %}
+  let pos = params.pastSeq + t;
+  for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = cos_cache[pos * HALF + d];
+    let sn = sin_cache[pos * HALF + d];
+    let x0 = k[d];
+    let x1 = k[d + HALF];
+    k[d] = x0 * cs - x1 * sn;
+    k[d + HALF] = x1 * cs + x0 * sn;
+  }
+{% endif %}
+{% if mode == "build" %}
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[dstBase + d] = k[d];
+    present_value[dstBase + d] = src_v[srcBase + d];
+  }
+{% else %}
+{% if mode == "append_quant" %}
+  let activeEnd = u32(seqlens_k[b]) + 1u;
+  let appendStart = activeEnd - params.seq;
+  let dstBase = ((b * KV_HEADS + hk) * params.totalSeq + appendStart + t) * PACKED;
+{% else %}
+  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * PACKED;
+{% endif %}
+{% if bits == 8 %}
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
+    present_key[dstBase + d] = i32(clamp(round(k[d] / kscale(d, hk)), QMIN, QMAX));
+    present_value[dstBase + d] = i32(clamp(round(src_v[srcBase + d] / vscale(d, hk)), QMIN, QMAX));
+  }
+{% else %}
+  for (var dd = 0u; dd < PACKED; dd = dd + 1u) {
+    let d0 = dd * 2u;
+    let d1 = dd * 2u + 1u;
+    let qk0 = i32(clamp(round(k[d0] / kscale(d0, hk)), QMIN, QMAX));
+    let qk1 = i32(clamp(round(k[d1] / kscale(d1, hk)), QMIN, QMAX));
+    let qv0 = i32(clamp(round(src_v[srcBase + d0] / vscale(d0, hk)), QMIN, QMAX));
+    let qv1 = i32(clamp(round(src_v[srcBase + d1] / vscale(d1, hk)), QMIN, QMAX));
+    present_key[dstBase + dd] = (u32(qk0 + 8) & 0xFu) | ((u32(qk1 + 8) & 0xFu) << 4u);
+    present_value[dstBase + dd] = (u32(qv0 + 8) & 0xFu) | ((u32(qv1 + 8) & 0xFu) << 4u);
+  }
+{% endif %}
+{% endif %}
+{% endif %}
+}
+{%- endif %}

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

	@@ -0,0 +1,57 @@

+{% if usesF16 is defined and usesF16 %}enable f16;
+{% endif %}{{ env.wgsl.resourceDeclarations }}
+// Q preprocessing for the flash-prefill path. One thread per (batch, qHead, qToken)
+// applies optional per-head qk-norm (RMS) then NeoX half-split rotary at the query's
+// absolute position, and writes a transformed Q (vec4, same BSH layout) that the
+// flash cluster then consumes unchanged. This mirrors the cooperative/threaded path's Q-prep
+// exactly (the present-cache K is already norm/rotary-applied, so only Q is touched),
+// so the flash result is bit-identical to the cooperative/threaded path.
+//
+// f16 queries are widened before cos/sin and norm/rotary arithmetic;
+// the qPrep output also stays f32 — RoPE precision matters at large positions, so
+// the rotated Q enters attention in full f32 (the cluster reads f32 Q + f16 K/V).
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
+const Q_HEADS: u32 = {{ qHeads }}u;
+const Q_HIDDEN: u32 = {{ qHidden }}u;
+const Q_HIDDEN_V4: u32 = {{ qHidden }}u / 4u;
+const WG: u32 = {{ copyWorkgroupSize }}u;
+const HALF: u32 = {{ half }}u;
+const QK_EPS: f32 = {{ qkEps }};
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the maxComputeWorkgroupsPerDimension
+  // workgroup-per-dimension dispatch limit. Reduces to gid.x when nwg.y == 1.
+  let qi = gid.x + gid.y * nwg.x * WG;
+  let total = params.batch * Q_HEADS * params.qSeq;
+  if (qi >= total) { return; }
+  let s = qi % params.qSeq;
+  let tmp = qi / params.qSeq;
+  let h = tmp % Q_HEADS;
+  let b = tmp / Q_HEADS;
+  let absPos = params.totalSeq - params.qSeq + s;
+  let base = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;
+  var q: array<f32, HEAD_DIM>;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = f32(query[base + d]); }
+  var ss = 0.0;
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ss = ss + q[d] * q[d]; }
+  let invRms = inverseSqrt(ss / f32(HEAD_DIM) + QK_EPS);
+  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = q[d] * invRms * f32(q_norm_weight[d]); }
+  for (var d = 0u; d < HALF; d = d + 1u) {
+    let cs = f32(cos_cache[absPos * HALF + d]);
+    let sn = f32(sin_cache[absPos * HALF + d]);
+    let x0 = q[d];
+    let x1 = q[d + HALF];
+    q[d] = x0 * cs - x1 * sn;
+    q[d + HALF] = x1 * cs + x0 * sn;
+  }
+  let base4 = (b * params.qSeq + s) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
+  for (var c = 0u; c < HEAD_DIM_V4; c = c + 1u) {
+    qout[base4 + c] = vec4<f32>(q[c * 4u], q[c * 4u + 1u], q[c * 4u + 2u], q[c * 4u + 3u]);
+  }
+}

build/webgpu/manifest.json ADDED Viewed

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

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,27 @@

+{
+  "name": "com.microsoft.GroupQueryAttention",
+  "id": "_com_microsoft_groupqueryattention_webgpu_75fbc9b",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "attention-rank4-tiled.wgsl.jinja": "EmYf/jV5lG4hmXr5wP+Qo6hZlY3LDd1aeAQhFXtDJEc=",
+      "attn-flash-decode-splitk-merge.wgsl.jinja": "kQDmryD3bdfkboFSLPK4pF2D40t3MX2xgJrMHx031yc=",
+      "attn-flash-decode-splitk.wgsl.jinja": "M6AWhYqX3pJ/MHz2K6Vy1rx7RjaZktzgQqB9GoY4cj0=",
+      "attn-flash-online.wgsl.jinja": "MStTMF1/l5L9Rzx1+rPz5T+70/46n/vO/+hZQ0lbFws=",
+      "attn-flash-prefill-cluster.wgsl.jinja": "yk1WymJIbxCy0slbiho5DS223uZH9PQ15629yvQ9RFc=",
+      "attn-flash-q32-broadcast.wgsl.jinja": "evSdaiYgvrikYZ8hpb5m4gxf+rIryKkHOTiKgyF7ioI=",
+      "attn-online-scalar.wgsl.jinja": "X4yGNCYu9+uJX6Q3EKkxV6v28CPlIYWme0Nmf8F16L8=",
+      "bench.json": "rGPC3cQInjJZjL+u+MzVXCOGu4L+G0v+/0UGDI9DEuQ=",
+      "gqa-attention.wgsl.jinja": "0EbwmdstNZMuskS170zYEo8whbQOFtupXLnNP2PaFL0=",
+      "gqa-present.wgsl.jinja": "J9g6oAeY2BtbmWBMlmbIXbxzo4sKHdxYqcE4EAd7BYM=",
+      "gqa-qprep.wgsl.jinja": "SvA19/AQoN2FXduwBtMI5QzBeeth2y9g8EnaIkV3V/0=",
+      "manifest.json": "qODDtl2Snem6J/K6QJwIL5i6CnGQ898t0a/EEka4JnY=",
+      "test.json": "iZAuQQINr2i9Q2hwXzNhKO+UT+ByfRwPsNc+paVhR20="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GroupQueryAttention" }
+}

build/webgpu/test.json ADDED Viewed

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