sync 6fdf6301e2bb

Browse files

Files changed (7) hide show

README.md +118 -0
build/webgpu/attn-packed-varlen-cluster.wgsl.jinja +432 -0
build/webgpu/attn-packed-varlen-scalar.wgsl.jinja +187 -0
build/webgpu/bench.json +0 -0
build/webgpu/manifest.json +667 -0
build/webgpu/metadata.json +47 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.PackedMultiHeadAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Multi-head self-attention over a padding-removed token stream. The `token_offset` and `cumulative_sequence_length` schedule maps every packed token to its sequence, and a token attends only the keys of that sequence. Query is either packed `[T, N, 3, H]` with key and value absent, or a separate `[T, hidden]` triple. The optional attention bias is indexed in PADDED coordinates. Float32 and float16 only; bfloat16 is not implemented.
+See the [ONNX Runtime `PackedMultiHeadAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.PackedMultiHeadAttention) for the reference semantics.
+## Inputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `queryT` | `query` | `T` | same as logical dtype | — | — | Packed queries: `[token_count, num_heads, 3, head_size]` with `key` and `value` absent, or `[token_count, hidden_size]` alongside both. | required |
+| `keyT` | `key` | `T` | same as logical dtype | `2` | — | Keys with shape `[token_count, hidden_size]`, matching `query` exactly. Present only in the separate form. | optional |
+| `valueT` | `value` | `T` | same as logical dtype | `2` | — | Values with shape `[token_count, v_hidden_size]`. Present only in the separate form. | optional |
+| `biasT` | `bias` | `T` | same as logical dtype | `1` | — | Optional packed input-projection bias with shape `[hidden_size + hidden_size + v_hidden_size]`, added to Q, K and V. | optional |
+| `tokenOffsetT` | `token_offset` | `M` | `int32` | `2` | — | Shape `[batch_size, sequence_length]`. The first `token_count` entries hold each packed token's flat index in the padded grid; the rest hold the padding positions. This tensor carries the padded `sequence_length`. | required |
+| `cumulativeSequenceLengthT` | `cumulative_sequence_length` | `M` | `int32` | `1` | — | Exclusive prefix sums with shape `[batch_size + 1]`, starting at 0 and ending at `token_count`. Sequence `i` owns packed tokens `[cum[i], cum[i + 1])`. Outputs are unspecified for a malformed schedule. | required |
+| `attentionBiasT` | `attention_bias` | `T` | same as logical dtype | `4` | — | Optional additive score bias with shape `[batch_size or 1, num_heads or 1, sequence_length, sequence_length]`. Its trailing axes are PADDED positions inside a sequence, not packed indices. | optional |
+## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | `2` | derived | Attention output with shape `[token_count, v_hidden_size]`, in packed token order. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `num_heads` | — | Number of attention heads. Required; the hidden size must be divisible by it. |
+| `scale` | — | Optional score scale. Omission, and the value `0`, both select `1 / sqrt(head_size)`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `packed_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `packed_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `packed_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `packed_attn_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `packed_attn_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `packed_attn_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `packed_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `packed_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `packed_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `packed_bias_attn_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `packed_bias_attn_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `packed_bias_attn_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `separate_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `separate_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `separate_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `separate_attn_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `separate_attn_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `separate_attn_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `separate_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `separate_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `separate_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `separate_bias_attn_bias_flash` — Tiled attention with independent Q/K and V slices and subgroup dot reductions. Equal wide heads target four `vec4` fragments per lane, bounded by subgroup/workgroup limits and native staging capacity. Unequal tiles cap work by token count and yield to serial or wider portable clusters when reuse is limited.
+- `separate_bias_attn_bias_flash_nosg` — Portable tiled attention with independent query/key and value slices. Lanes cover the larger width, with guarded zero Q/K fragments and shared reductions. Device limits bound lane count and staging. Unequal-width tiles yield to serial attention below the query-reuse target.
+- `separate_bias_attn_bias_serial` — Portable serial-key attention: one workgroup owns one (packed token, head), walks that token's own sequence and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+## Device requirements
+Some implementation variants require `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`attn-packed-varlen-cluster.wgsl.jinja`](build/webgpu/attn-packed-varlen-cluster.wgsl.jinja)
+- [`attn-packed-varlen-scalar.wgsl.jinja`](build/webgpu/attn-packed-varlen-scalar.wgsl.jinja)
+## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.PackedMultiHeadAttention", { version: 1 });
+const { outputT } = await kernel({
+  queryT: { data: queryTData, shape: [1, 16] },
+  keyT: { data: keyTData, shape: [1, 16] },
+  valueT: { data: valueTData, shape: [1, 20] },
+  tokenOffsetT: { data: tokenOffsetTData, shape: [1, 1] },
+  cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [2] },
+}, {
+  attrs: { num_heads: 1 },
+});
+```

build/webgpu/attn-packed-varlen-cluster.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,432 @@

+{% set parallelSoftmax = distributedSoftmax is defined and distributedSoftmax and not useSubgroups %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set QBUF = "query" %}
+{% set KBUF = "query" if packedQkv else "key" %}
+{% set VBUF = "query" if packedQkv else "value" %}
+{% set BIAS_K_ROW = biasKOffset if biasKOffset is defined else 0 %}
+{% set BIAS_V_ROW = biasVOffset if biasVOffset is defined else 0 %}
+{% set ST = "f16" if packedStageBytes == 2 else "f32" %}
+{% set SLICE_COUNT = (((headDimV4 + LPQ - 1) / LPQ) | int) %}
+{% set queryTail = headDimV4 % LPQ != 0 %}
+{% set VALUE_DIM_V4 = ((vHeadDim / 4) | int) %}
+{% set VALUE_SLICES = (((VALUE_DIM_V4 + LPQ - 1) / LPQ) | int) %}
+{% set unequalValue = vHeadDim != headDim %}
+{% set valueTail = VALUE_DIM_V4 % LPQ != 0 %}
+{% set valueWidth = "VALUE_HEAD_DIM_V4" if unequalValue else "HEAD_DIM_V4" %}
+{% set valueLane = "valueLaneSlice" if unequalValue else "laneSlice" %}
+// Tiled flash attention over a padding-removed packed token stream. One
+// workgroup covers TILE_Q consecutive packed tokens for one head, with an
+// LPQ-lane cluster per token; K and V are staged in workgroup memory once per
+// TILE_K keys and reused by every token in the tile. Each lane owns contiguous
+// query and output vec4 slices, sized independently for the Q/K and V widths.
+// Dot products, online softmax and weighted value sums accumulate in f32; output
+// conversion happens at the store. The staging type above governs biased K/V.
+// A packed tile can straddle a sequence boundary, so the key loop runs over the
+// workgroup's union range [wgStart, wgEnd) — the first token's sequence start to
+// the last token's sequence end — so the trip count is uniform, and each token
+// masks the keys outside its own sequence.
+{% if packedQkv %}
+// Source: one packed `[T, N, 3, H]` buffer; Q, K and V differ only by offset.
+{% else %}
+// Source: separate token-major `[T, hidden]` query, key and value buffers.
+{% endif %}
+{% if hasBias %}
+// Packed `[Q; K; V]` bias rows: the Q bias enters the query registers, and the
+// K and V biases fold into the staged tiles, so a bias element is added once per
+// tile rather than once per (token, key).
+{% endif %}
+{% if hasAttnBias %}
+// The additive attention bias lives in PADDED coordinates: its trailing two axes
+// are indexed by the query's and the key's position inside their shared
+// sequence, not by their packed indices. A zero broadcast stride collapses the
+// batch or head axis.
+{% endif %}
+{% if useSubgroups %}
+// The per-lane dot fragments combine through a log2(LPQ) subgroupShuffleXor
+// butterfly inside each cluster.
+{% else %}
+// The per-lane dot fragments combine through workgroup memory; this form
+// requires no subgroup feature.
+{% endif %}
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
+{% if unequalValue %}
+const VALUE_HEAD_DIM_V4: u32 = {{ VALUE_DIM_V4 }}u;
+{% if hasBias %}
+const VALUE_HEAD_DIM: u32 = {{ vHeadDim }}u;
+{% endif %}
+{% endif %}
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const Q_ROW_V4: u32 = {{ qRowStrideV4 }}u;
+const Q_HEAD_V4: u32 = {{ qHeadStrideV4 }}u;
+const K_ROW_V4: u32 = {{ kRowStrideV4 }}u;
+const K_HEAD_V4: u32 = {{ kHeadStrideV4 }}u;
+const K_OFFSET_V4: u32 = {{ kOffsetV4 }}u;
+const V_ROW_V4: u32 = {{ vRowStrideV4 }}u;
+const V_HEAD_V4: u32 = {{ vHeadStrideV4 }}u;
+const V_OFFSET_V4: u32 = {{ vOffsetV4 }}u;
+const O_ROW_V4: u32 = {{ oRowStrideV4 }}u;
+const O_HEAD_V4: u32 = {{ oHeadStrideV4 }}u;
+{% if hasBias %}
+const BIAS_K: u32 = {{ BIAS_K_ROW }}u;
+const BIAS_V: u32 = {{ BIAS_V_ROW }}u;
+{% endif %}
+const TILE_Q: u32 = {{ TILE_Q }}u;
+const TILE_K: u32 = {{ TILE_K }}u;
+// Lanes per token cluster, sized by the manifest for query and value ownership.
+const LPQ: u32 = {{ LPQ }}u;
+const SLICE: u32 = {% if queryTail %}{{ SLICE_COUNT }}u{% else %}HEAD_DIM_V4 / LPQ{% endif %};  // query vec4s per lane slice
+const WG: u32 = TILE_Q * LPQ;
+const NEG_INF: f32 = -3.4028234663852886e38;
+var<workgroup> k_tile: array<vec4<{{ ST }}>, TILE_K * HEAD_DIM_V4>;
+var<workgroup> v_tile: array<vec4<{{ ST }}>, TILE_K * {{ valueWidth }}>;
+{% if not useSubgroups %}
+// Reuse scratch for dot fragments, maxima, sums and probabilities when the
+// manifest admits the full tile; otherwise retain one key's reduction slots.
+var<workgroup> red: array<f32, {% if parallelSoftmax %}{{ TILE_Q * (LPQ * TILE_K if LPQ * TILE_K > 2 * LPQ + TILE_K else 2 * LPQ + TILE_K) }}{% else %}WG{% endif %}>;
+{% endif %}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+// Packed-token metadata. `token_offset` records, for every packed token, its
+// flat index in the padded `[batch, sequence]` grid, so the sequence index and
+// the position inside that sequence are one load and two integer operations.
+// Packing keeps a sequence's tokens contiguous and in ascending position order,
+// so the packed index of a sequence's first token is `token - position` and the
+// sequence's packed key range is [keyStart, keyEnd).
+struct PackedToken {
+  sequence: u32,
+  position: u32,
+  keyStart: u32,
+  keyEnd: u32,
+}
+fn packed_token_of(token: u32, seqPad: u32) -> PackedToken {
+  let flat = u32(token_offset[token]);
+  var info: PackedToken;
+  info.sequence = flat / seqPad;
+  info.position = flat - info.sequence * seqPad;
+  info.keyStart = token - info.position;
+  info.keyEnd = u32(cumulative_sequence_length[info.sequence + 1u]);
+  return info;
+}
+{% if hasBias %}
+{% set biasScalarWiden = usesF16 %}
+{% set BW = "f32(" if biasScalarWiden is defined and biasScalarWiden else "" %}
+{% set BC = ")" if biasScalarWiden is defined and biasScalarWiden else "" %}
+fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
+  let offset = base + d4 * 4u;
+  return vec4<f32>({{ BW }}bias[offset]{{ BC }}, {{ BW }}bias[offset + 1u]{{ BC }}, {{ BW }}bias[offset + 2u]{{ BC }}, {{ BW }}bias[offset + 3u]{{ BC }});
+}
+{% endif %}
+{% macro packed_dot_fragment() %}
+      var part: f32 = 0.0;
+      let kb = kk * HEAD_DIM_V4 + laneSlice;
+{% for c in range(SLICE_COUNT) %}
+{% if queryTail %}
+      if (laneSlice + {{ c }}u < HEAD_DIM_V4) {
+{% endif %}
+      part = part + dot(qr{{ c }}, vec4<f32>(k_tile[kb + {{ c }}u]));
+{% if queryTail %}
+      }
+{% endif %}
+{% endfor %}
+{% endmacro %}
+{% macro packed_value_accumulate() %}
+        let vRow = kk * {{ valueWidth }} + {{ valueLane }};
+{% for c in range(VALUE_SLICES) %}
+{% if valueTail %}
+        if ({{ valueLane }} + {{ c }}u < {{ valueWidth }}) {
+{% endif %}
+        o{{ c }} = o{{ c }} + p * vec4<f32>(v_tile[vRow + {{ c }}u]);
+{% if valueTail %}
+        }
+{% endif %}
+{% endfor %}
+{% endmacro %}
+{% macro packed_score_write(score_index = "kk") %}
+      if (kj >= minKj && kj < maxKj) {
+{% if hasAttnBias %}
+        let maskIndex = info.sequence * params.maskBatchStride + h * params.maskHeadStride
+          + info.position * params.maskSeqStride + (kj - minKj);
+        s[{{ score_index }}] = part * SCALE + f32(attn_bias[maskIndex]);
+{% else %}
+        s[{{ score_index }}] = part * SCALE;
+{% endif %}
+      }
+{% endmacro %}
+@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;
+  // The query-tile axis is folded over x and z with a baked stride; z is 1
+  // whenever the tile count fits one dimension, so the fold term vanishes there.
+  let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
+  let firstToken = tile * TILE_Q;
+  if (firstToken >= params.tokenCount || h >= Q_HEADS) {
+    return;
+  }
+{% if useSubgroups %}
+  // The subgroupShuffleXor butterfly requires every LPQ-lane 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 tid = lid.x;
+  let qSub = tid / LPQ;
+  let lane = tid % LPQ;
+  let qIdx = firstToken + qSub;
+  let qValid = qIdx < params.tokenCount;
+  let qClamped = min(qIdx, params.tokenCount - 1u);
+  let info = packed_token_of(qClamped, params.seqPad);
+  let minKj = info.keyStart;
+  let maxKj = info.keyEnd;
+  // Union key range for the whole tile: packed tokens are ordered, so the first
+  // token's sequence start and the last token's sequence end bound every token
+  // in the tile.
+  let lastToken = min(firstToken + TILE_Q - 1u, params.tokenCount - 1u);
+  let wgStart = packed_token_of(firstToken, params.seqPad).keyStart;
+  let wgEnd = packed_token_of(lastToken, params.seqPad).keyEnd;
+  let SCALE = scale_value();
+  let laneSlice = lane * SLICE;
+{% if unequalValue %}
+  let valueLaneSlice = lane * {{ VALUE_SLICES }}u;
+{% endif %}
+  let qBase4 = qClamped * Q_ROW_V4 + h * Q_HEAD_V4 + laneSlice;
+{% for c in range(SLICE_COUNT) %}
+{% if queryTail %}
+  var qr{{ c }} = vec4<f32>(0.0);
+  if (laneSlice + {{ c }}u < HEAD_DIM_V4) {
+    qr{{ c }} = vec4<f32>({{ QBUF }}[qBase4 + {{ c }}u]);
+{% if hasBias %}
+    qr{{ c }} = qr{{ c }} + load_bias4(h * HEAD_DIM, laneSlice + {{ c }}u);
+{% endif %}
+  }
+{% else %}
+  var qr{{ c }} = vec4<f32>({{ QBUF }}[qBase4 + {{ c }}u]);
+{% if hasBias %}
+  qr{{ c }} = qr{{ c }} + load_bias4(h * HEAD_DIM, laneSlice + {{ c }}u);
+{% endif %}
+{% endif %}
+{% if not unequalValue %}
+  var o{{ c }} = vec4<f32>(0.0);
+{% endif %}
+{% endfor %}
+{% if unequalValue %}
+{% for c in range(VALUE_SLICES) %}
+  var o{{ c }} = vec4<f32>(0.0);
+{% endfor %}
+{% endif %}
+  var m: f32 = NEG_INF;
+  var l: f32 = 0.0;
+  var kStart: u32 = wgStart;
+  loop {
+    if (kStart >= wgEnd) { break; }
+    // Cooperative K/V tile load with out-of-range keys zero-filled. The K and V
+    // biases are token-independent, so they fold in here, once per staged word.
+    workgroupBarrier();
+{% if unequalValue %}
+{% for tile, buffer, rowStride, headStride, offset, width, biasRow, biasWidth in [
+  ["k_tile", KBUF, "K_ROW_V4", "K_HEAD_V4", "K_OFFSET_V4", "HEAD_DIM_V4", "BIAS_K", "HEAD_DIM"],
+  ["v_tile", VBUF, "V_ROW_V4", "V_HEAD_V4", "V_OFFSET_V4", "VALUE_HEAD_DIM_V4", "BIAS_V", "VALUE_HEAD_DIM"]
+] %}
+    for (var i: u32 = tid; i < TILE_K * {{ width }}; i = i + WG) {
+      let slot = i / {{ width }};
+      let d4 = i % {{ width }};
+      let kj = kStart + slot;
+      if (kj < wgEnd) {
+        var word = vec4<f32>({{ buffer }}[kj * {{ rowStride }} + h * {{ headStride }} + {{ offset }} + d4]);
+{% if hasBias %}
+        word = word + load_bias4({{ biasRow }} + h * {{ biasWidth }}, d4);
+{% endif %}
+        {{ tile }}[i] = vec4<{{ ST }}>(word);
+      } else {
+        {{ tile }}[i] = vec4<{{ ST }}>(0.0);
+      }
+    }
+{% endfor %}
+{% else %}
+    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) {
+        var kWord = vec4<f32>({{ KBUF }}[kj * K_ROW_V4 + h * K_HEAD_V4 + K_OFFSET_V4 + d4]);
+        var vWord = vec4<f32>({{ VBUF }}[kj * V_ROW_V4 + h * V_HEAD_V4 + V_OFFSET_V4 + d4]);
+{% if hasBias %}
+        kWord = kWord + load_bias4(BIAS_K + h * HEAD_DIM, d4);
+        vWord = vWord + load_bias4(BIAS_V + h * HEAD_DIM, d4);
+{% endif %}
+        k_tile[i] = vec4<{{ ST }}>(kWord);
+        v_tile[i] = vec4<{{ ST }}>(vWord);
+      } else {
+        k_tile[i] = vec4<{{ ST }}>(0.0);
+        v_tile[i] = vec4<{{ ST }}>(0.0);
+      }
+    }
+{% endif %}
+    workgroupBarrier();
+{% if parallelSoftmax %}
+    // Each cluster lane owns disjoint key scores and exponentials. The dot
+    // fragments, maxima, sums and probabilities reuse one shared allocation.
+    var s: array<f32, {{ ((TILE_K + LPQ - 1) / LPQ) | int }}>;
+    for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+{{ packed_dot_fragment() }}
+      red[kk * WG + tid] = part;
+    }
+    workgroupBarrier();
+    var localMax = NEG_INF;
+    for (var kk: u32 = lane; kk < TILE_K; kk = kk + LPQ) {
+      let kj = kStart + kk;
+      let key_slot = kk / LPQ;
+      s[key_slot] = NEG_INF;
+      var part: f32 = 0.0;
+      let cbase = kk * WG + tid - lane;
+      for (var j: u32 = 0u; j < LPQ; j = j + 1u) {
+        part = part + red[cbase + j];
+      }
+{{ packed_score_write("key_slot") }}
+      localMax = max(localMax, s[key_slot]);
+    }
+    // All fragment readers finish before their storage is reused for maxima.
+    workgroupBarrier();
+    red[tid] = localMax;
+    workgroupBarrier();
+    let cbase = tid - lane;
+    // Keyless lanes contribute only NEG_INF or zero to max/sum. Dot fragments
+    // above still consume every lane, including when TILE_K is smaller than LPQ.
+    var tileMax = NEG_INF;
+    for (var j: u32 = 0u; j < {{ "LPQ" if TILE_K >= LPQ else "TILE_K" }}; j = j + 1u) {
+      tileMax = max(tileMax, red[cbase + j]);
+    }
+    let newMax = max(m, tileMax);
+    let corr = select(exp(m - newMax), 0.0, m == NEG_INF);
+    var localSum: f32 = 0.0;
+    for (var kk: u32 = lane; kk < TILE_K; kk = kk + LPQ) {
+      let score = s[kk / LPQ];
+      let pk = select(0.0, exp(score - newMax), score != NEG_INF);
+      red[2u * WG + qSub * TILE_K + kk] = pk;
+      localSum = localSum + pk;
+    }
+    // Sum slots are disjoint from the maxima still being consumed above.
+    red[WG + tid] = localSum;
+    workgroupBarrier();
+    var pSum: f32 = 0.0;
+    for (var j: u32 = 0u; j < {{ "LPQ" if TILE_K >= LPQ else "TILE_K" }}; j = j + 1u) {
+      pSum = pSum + red[WG + cbase + j];
+    }
+    l = l * corr + pSum;
+    m = newMax;
+{% 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;
+{{ packed_dot_fragment() }}
+{% if useSubgroups %}
+      // log2(LPQ) butterfly all-reduce inside the cluster: every lane ends with
+      // the full dot.
+      part = part + subgroupShuffleXor(part, 1u);
+{% for shift in [2, 4, 8, 16, 32] %}
+{% if LPQ > shift %}
+      part = part + subgroupShuffleXor(part, {{ shift }}u);
+{% endif %}
+{% endfor %}
+{% else %}
+      red[tid] = part;
+      workgroupBarrier();
+      let cbase = tid - lane;
+      part = 0.0;
+      for (var j: u32 = 0u; j < LPQ; j = j + 1u) {
+        part = part + red[cbase + j];
+      }
+      workgroupBarrier();
+{% endif %}
+{{ packed_score_write() }}
+    }
+    // Per-lane online softmax over the tile. s[kk] is reused to hold the
+    // exponentiated probabilities for the value 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;
+{% endif %}
+    // A boundary tile addresses V rows outside this token's sequence. Because
+    // 0 * NaN is NaN, the guarded loop skips the V read entirely for a key with
+    // no probability mass; an interior tile runs the unguarded FMA chain. One
+    // range test covers the whole lane slice, and one probability load feeds
+    // every vec4 of it.
+    let tileInterior = kStart >= minKj && kStart + TILE_K <= maxKj;
+{% for c in range(VALUE_SLICES) %}
+    o{{ c }} = o{{ c }} * corr;
+{% endfor %}
+    if (tileInterior) {
+      for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+        let p = {% if parallelSoftmax %}red[2u * WG + qSub * TILE_K + kk]{% else %}s[kk]{% endif %};
+{{ packed_value_accumulate() }}
+      }
+    } else {
+      for (var kk: u32 = 0u; kk < TILE_K; kk = kk + 1u) {
+        let p = {% if parallelSoftmax %}red[2u * WG + qSub * TILE_K + kk]{% else %}s[kk]{% endif %};
+        if (p != 0.0) {
+{{ packed_value_accumulate() }}
+        }
+      }
+    }
+    kStart = kStart + TILE_K;
+  }
+  if (qValid) {
+    // l == 0 means this token had no probability-bearing key. Emit 0 rather
+    // than 0/0.
+    let inv = select(0.0, 1.0 / l, l > 0.0);
+    let outBase4 = qIdx * O_ROW_V4 + h * O_HEAD_V4 + {{ valueLane }};
+{% for c in range(VALUE_SLICES) %}
+{% if valueTail %}
+    if ({{ valueLane }} + {{ c }}u < {{ valueWidth }}) {
+{% endif %}
+    output[outBase4 + {{ c }}u] = vec4<{{ scalar }}>(o{{ c }} * inv);
+{% if valueTail %}
+    }
+{% endif %}
+{% endfor %}
+  }
+}

build/webgpu/attn-packed-varlen-scalar.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,187 @@

+{{ env.wgsl.resourceDeclarations }}
+{% set QBUF = "query" %}
+{% set KBUF = "query" if packedQkv else "key" %}
+{% set VBUF = "query" if packedQkv else "value" %}
+{% set BIAS_K_ROW = biasKOffset if biasKOffset is defined else 0 %}
+{% set BIAS_V_ROW = biasVOffset if biasVOffset is defined else 0 %}
+// Online-softmax attention over a padding-removed packed token stream, with no
+// feature requirements. One workgroup owns one (packed token, query head): the
+// workgroup walks its own sequence's keys serially, cooperates on each q·k dot
+// through a shared-memory tree reduction, and applies the online rescale to a
+// running value accumulator per key. Head dimensions need no alignment and the
+// V head dimension may differ from the Q/K head dimension.
+{% if packedQkv %}
+// Source: one packed `[T, N, 3, H]` buffer; Q, K and V differ only by offset.
+{% else %}
+// Source: separate token-major `[T, hidden]` query, key and value buffers.
+{% endif %}
+{% if hasBias %}
+// Packed `[Q; K; V]` bias rows are added to Q, K and V during the key walk,
+// which is where the reference adds them.
+{% endif %}
+{% if hasAttnBias %}
+// The additive attention bias lives in PADDED coordinates: its trailing two
+// axes are indexed by the query's and the key's position inside their shared
+// sequence, not by their packed indices. A zero broadcast stride collapses the
+// batch or head axis.
+{% endif %}
+const HEAD_DIM: u32 = {{ headDim }}u;
+const V_HEAD_DIM: u32 = {{ vHeadDim }}u;
+const Q_HEADS: u32 = {{ qNumHeads }}u;
+const Q_ROW: u32 = {{ qRowStride }}u;
+const Q_HEAD: u32 = {{ qHeadStride }}u;
+const K_ROW: u32 = {{ kRowStride }}u;
+const K_HEAD: u32 = {{ kHeadStride }}u;
+const K_OFFSET: u32 = {{ kOffset }}u;
+const V_ROW: u32 = {{ vRowStride }}u;
+const V_HEAD: u32 = {{ vHeadStride }}u;
+const V_OFFSET: u32 = {{ vOffset }}u;
+const O_ROW: u32 = {{ oRowStride }}u;
+const O_HEAD: u32 = {{ oHeadStride }}u;
+{% if hasBias %}
+const BIAS_K: u32 = {{ BIAS_K_ROW }}u;
+const BIAS_V: u32 = {{ BIAS_V_ROW }}u;
+{% endif %}
+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, V_HEAD_DIM>;
+var<workgroup> previous_scale: f32;
+{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %}
+fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} {
+  {{ buffer }}[tid] = value;
+  workgroupBarrier();
+  // Ceil-halving keeps every lane when the workgroup size is not a power of
+  // 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) {
+      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
+    }
+    workgroupBarrier();
+    n = half;
+    if (n == 1u) {
+      break;
+    }
+  }
+  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
+  // slot 0 here, so the next call's first store must not run until all lanes have read it.
+  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
+  let reduced = {{ buffer }}[0];
+  workgroupBarrier();
+  return reduced;
+}
+{% endmacro %}
+{{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
+{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
+  if (params.scale != 0.0) { return params.scale; }
+  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
+}
+// Packed-token metadata. `token_offset` records, for every packed token, its
+// flat index in the padded `[batch, sequence]` grid, so the sequence index and
+// the position inside that sequence are one load and two integer operations.
+// Packing keeps a sequence's tokens contiguous and in ascending position order,
+// so the packed index of a sequence's first token is `token - position` and the
+// sequence's packed key range is [keyStart, keyEnd).
+struct PackedToken {
+  sequence: u32,
+  position: u32,
+  keyStart: u32,
+  keyEnd: u32,
+}
+fn packed_token_of(token: u32, seqPad: u32) -> PackedToken {
+  let flat = u32(token_offset[token]);
+  var info: PackedToken;
+  info.sequence = flat / seqPad;
+  info.position = flat - info.sequence * seqPad;
+  info.keyStart = token - info.position;
+  info.keyEnd = u32(cumulative_sequence_length[info.sequence + 1u]);
+  return info;
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  // The token axis is folded over x and z with a baked stride; z is 1 whenever
+  // the token count fits one dimension, so the fold term vanishes there.
+  let token = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
+  let h = wg.y;
+  if (token >= params.tokenCount || h >= Q_HEADS) {
+    return;
+  }
+  let tid = lid.x;
+  let info = packed_token_of(token, params.seqPad);
+  let qBase = token * Q_ROW + h * Q_HEAD;
+  let oBase = token * O_ROW + h * O_HEAD;
+  if (tid == 0u) {
+    running_max = -3.4028234663852886e38;
+    running_denom = 0.0;
+  }
+  for (var d: u32 = tid; d < V_HEAD_DIM; d = d + WG) {
+    running_out[d] = 0.0;
+  }
+  workgroupBarrier();
+  for (var keyToken: u32 = info.keyStart; keyToken < info.keyEnd; keyToken = keyToken + 1u) {
+    let kBase = keyToken * K_ROW + h * K_HEAD + K_OFFSET;
+    var partial_dot = 0.0;
+    for (var d: u32 = tid; d < HEAD_DIM; d = d + WG) {
+      var q_value = f32({{ QBUF }}[qBase + d]);
+      var k_value = f32({{ KBUF }}[kBase + d]);
+{% if hasBias %}
+      let channel = h * HEAD_DIM + d;
+      q_value = q_value + f32(bias[channel]);
+      k_value = k_value + f32(bias[BIAS_K + 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/denominator
+    // (broadcast through shared memory by the barrier below).
+{% if hasAttnBias %}
+    let maskIndex = info.sequence * params.maskBatchStride + h * params.maskHeadStride
+      + info.position * params.maskSeqStride + (keyToken - info.keyStart);
+    let score = reduce_sum(partial_dot, tid) * scale_value() + f32(attn_bias[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);
+    let vBase = keyToken * V_ROW + h * V_HEAD + V_OFFSET;
+    for (var d: u32 = tid; d < V_HEAD_DIM; d = d + WG) {
+      var v_value = f32({{ VBUF }}[vBase + d]);
+{% if hasBias %}
+      v_value = v_value + f32(bias[BIAS_V + h * V_HEAD_DIM + d]);
+{% endif %}
+      running_out[d] = running_out[d] * previous_scale + probability_numerator * v_value;
+    }
+    workgroupBarrier();
+  }
+  // A token whose sequence carries no key leaves running_denom at zero; emit a
+  // zero row rather than 0 * (1/0) = NaN.
+  let inv_denom = select(0.0, 1.0 / running_denom, running_denom > 0.0);
+  for (var d: u32 = tid; d < V_HEAD_DIM; d = d + WG) {
+    output[oBase + d] = {{ outputScalar }}(running_out[d] * inv_denom);
+  }
+}

build/webgpu/bench.json ADDED Viewed

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

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,667 @@

+{
+  "domain": "com.microsoft",
+  "name": "PackedMultiHeadAttention",
+  "sinceVersion": 1,
+  "inputs": {
+    "queryT": { "onnx": "query", "dtype": "T" },
+    "keyT": { "onnx": "key", "dtype": "T", "rank": 2, "optional": true },
+    "valueT": { "onnx": "value", "dtype": "T", "rank": 2, "optional": true },
+    "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true },
+    "tokenOffsetT": { "onnx": "token_offset", "dtype": "M", "rank": 2, "storage": "int32" },
+    "cumulativeSequenceLengthT": { "onnx": "cumulative_sequence_length", "dtype": "M", "rank": 1, "storage": "int32" },
+    "attentionBiasT": { "onnx": "attention_bias", "dtype": "T", "rank": 4, "optional": true }
+  },
+  "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 2, "shape": "[tokenCount, vHiddenSize]" } },
+  "attributes": { "num_heads": {}, "scale": {} },
+  "attributeConstraints": { "num_heads": { "required": true } },
+  "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
+  "tunables": { "CLUSTER_TILE_K": { "default": 16 }, "SCALAR_WORKGROUP_SIZE": { "default": 64 } },
+  "derive": {
+    "numHeads": "attrs.num_heads",
+    "qNumHeads": "numHeads",
+    "packedQkv": "ranks.queryT == 4",
+    "tokenCount": "dim(shapes.queryT, 0)",
+    "batchSize": "dim(shapes.tokenOffsetT, 0) if ranks.tokenOffsetT == 2 else 0",
+    "seqPad": "dim(shapes.tokenOffsetT, 1) if ranks.tokenOffsetT == 2 else 0",
+    "packedHeadDim": "dim(shapes.queryT, 3) if packedQkv else 0",
+    "hiddenSize": "numHeads * packedHeadDim if packedQkv else (dim(shapes.queryT, 1) if ranks.queryT == 2 else 0)",
+    "vHiddenSize": "hiddenSize if packedQkv else (dim(shapes.valueT, 1) if present.valueT and ranks.valueT == 2 else 0)",
+    "headDim": "hiddenSize / numHeads if numHeads > 0 and hiddenSize % numHeads == 0 else 0",
+    "vHeadDim": "vHiddenSize / numHeads if numHeads > 0 and vHiddenSize % numHeads == 0 else 0",
+    "attentionScale": "attrs.scale if has(attrs, \"scale\") else 0",
+    "qRowStride": "numHeads * 3 * headDim if packedQkv else hiddenSize",
+    "qHeadStride": "3 * headDim if packedQkv else headDim",
+    "kRowStride": "qRowStride if packedQkv else hiddenSize",
+    "kHeadStride": "qHeadStride if packedQkv else headDim",
+    "kOffset": "headDim if packedQkv else 0",
+    "vRowStride": "qRowStride if packedQkv else vHiddenSize",
+    "vHeadStride": "qHeadStride if packedQkv else vHeadDim",
+    "vOffset": "2 * headDim if packedQkv else 0",
+    "oRowStride": "vHiddenSize",
+    "oHeadStride": "vHeadDim",
+    "biasKOffset": "hiddenSize",
+    "biasVOffset": "2 * hiddenSize",
+    "attnBiasHeads": "dim(shapes.attentionBiasT, 1) if present.attentionBiasT and ranks.attentionBiasT == 4 else 0",
+    "maskSeqStride": "seqPad",
+    "maskHeadStride": "0 if attnBiasHeads == 1 else seqPad * seqPad",
+    "maskBatchStride": "0 if (not present.attentionBiasT) or ranks.attentionBiasT != 4 or dim(shapes.attentionBiasT, 0) == 1 else attnBiasHeads * seqPad * seqPad",
+    "dtypeOk": "tensorDtypes.outputT == tensorDtypes.queryT and (tensorDtypes.queryT == \"float32\" or tensorDtypes.queryT == \"float16\") and f16Ok(dtypes.T)",
+    "packedShapeOk": "packedQkv and not present.keyT and not present.valueT and dim(shapes.queryT, 1) == numHeads and dim(shapes.queryT, 2) == 3 and packedHeadDim > 0",
+    "separateShapeOk": "ranks.queryT == 2 and present.keyT and present.valueT and ranks.keyT == 2 and ranks.valueT == 2 and sameShape(shapes.keyT, shapes.queryT) and dim(shapes.valueT, 0) == tokenCount and tensorDtypes.keyT == tensorDtypes.queryT and tensorDtypes.valueT == tensorDtypes.queryT",
+    "scheduleOk": "ranks.tokenOffsetT == 2 and ranks.cumulativeSequenceLengthT == 1 and dim(shapes.cumulativeSequenceLengthT, 0) == batchSize + 1 and batchSize > 0 and seqPad > 0 and numel(shapes.tokenOffsetT) >= tokenCount",
+    "outputShapeOk": "ranks.outputT == 2 and dim(shapes.outputT, 0) == tokenCount and dim(shapes.outputT, 1) == vHiddenSize",
+    "biasOk": "present.biasT and ranks.biasT == 1 and dim(shapes.biasT, 0) == 2 * hiddenSize + vHiddenSize and tensorDtypes.biasT == tensorDtypes.queryT",
+    "attnBiasOk": "present.attentionBiasT and ranks.attentionBiasT == 4 and (dim(shapes.attentionBiasT, 0) == 1 or dim(shapes.attentionBiasT, 0) == batchSize) and (attnBiasHeads == 1 or attnBiasHeads == numHeads) and dim(shapes.attentionBiasT, 2) == seqPad and dim(shapes.attentionBiasT, 3) == seqPad and tensorDtypes.attentionBiasT == tensorDtypes.queryT",
+    "hasBias": "present.biasT",
+    "hasAttnBias": "present.attentionBiasT",
+    "baseContract": "dtypeOk and scheduleOk and outputShapeOk and numHeads > 0 and headDim > 0 and vHeadDim > 0 and tokenCount > 0 and (packedShapeOk or separateShapeOk)",
+    "usesF16": "tensorDtypes.queryT == \"float16\"",
+    "scalar": "\"f16\" if usesF16 else \"f32\"",
+    "outputScalar": "scalar",
+    "inputVec4": "\"vec4<f16>\" if usesF16 else \"vec4<f32>\"",
+    "qkvBytes": "2 if usesF16 else 4",
+    "packedStageBytes": "4 if usesF16 and present.biasT and vHeadDim != headDim else qkvBytes",
+    "clusterWidthAligned": "headDim % 16 == 0 and vHeadDim % 4 == 0",
+    "headDimV4": "headDim / 4",
+    "qRowStrideV4": "qRowStride / 4",
+    "qHeadStrideV4": "qHeadStride / 4",
+    "kRowStrideV4": "kRowStride / 4",
+    "kHeadStrideV4": "kHeadStride / 4",
+    "kOffsetV4": "kOffset / 4",
+    "vRowStrideV4": "vRowStride / 4",
+    "vHeadStrideV4": "vHeadStride / 4",
+    "vOffsetV4": "vOffset / 4",
+    "oRowStrideV4": "oRowStride / 4",
+    "oHeadStrideV4": "oHeadStride / 4",
+    "scalarWg": "tunables.SCALAR_WORKGROUP_SIZE",
+    "scalarBytes": "4 * scalarWg + 4 * vHeadDim + 12",
+    "scalarOk": "scalarWg <= device.limits.maxComputeInvocationsPerWorkgroup and scalarWg <= device.limits.maxComputeWorkgroupSizeX and scalarBytes <= device.limits.maxComputeWorkgroupStorageSize and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "clusterTileQ": "min(16, pow2ceil(tokenCount)) if vHeadDim != headDim else 16",
+    "clusterBaseLpq": "8 if headDim % 32 == 0 else 4",
+    "clusterTileK": "min(tunables.CLUSTER_TILE_K, pow2ceil(tokenCount)) if vHeadDim != headDim else tunables.CLUSTER_TILE_K",
+    "clusterTileBytes": "clusterTileK * (headDim + vHeadDim) * packedStageBytes",
+    "clusterLaneTarget": "min(32, max(clusterBaseLpq, pow2ceil(headDim / 16)), device.adapterInfo.subgroupMinSize, pow2ceil(floor(min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX) / clusterTileQ) + 1) / 2) if vHeadDim == headDim and clusterWidthAligned and clusterTileBytes <= device.limits.maxComputeWorkgroupStorageSize and device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and device.adapterInfo.subgroupMinSize >= clusterBaseLpq else clusterBaseLpq",
+    "clusterLpq": "clusterLaneTarget if ceilDiv(headDim, 4 * clusterBaseLpq) >= 2 * ceilDiv(headDim, 4 * clusterLaneTarget) else clusterBaseLpq",
+    "clusterWg": "clusterTileQ * clusterLpq",
+    "clusterNosgBytes": "clusterTileBytes + 4 * clusterWg",
+    "clusterTiles": "ceilDiv(tokenCount, clusterTileQ)",
+    "clusterShapeOk": "clusterWidthAligned and clusterWg <= device.limits.maxComputeInvocationsPerWorkgroup and clusterWg <= device.limits.maxComputeWorkgroupSizeX and numHeads <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "clusterOk": "clusterShapeOk and clusterTileBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "clusterNosgOk": "clusterShapeOk and clusterNosgBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "clusterSubgroupsOk": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize >= clusterLpq and device.adapterInfo.subgroupMinSize % clusterLpq == 0 and device.adapterInfo.subgroupMaxSize >= device.adapterInfo.subgroupMinSize and device.adapterInfo.subgroupMaxSize % clusterLpq == 0",
+    "portableLaneTarget": "min(64, pow2ceil(max(headDimV4, vHeadDim / 4) / 4), device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "portableLpq": "max(clusterLpq, pow2ceil(portableLaneTarget + 1) / 2) if vHeadDim != headDim else (64 if portableLaneTarget >= 64 and headDimV4 % 64 == 0 else (32 if portableLaneTarget >= 32 and headDimV4 % 32 == 0 else (16 if portableLaneTarget >= 16 and headDimV4 % 16 == 0 else clusterLpq)))",
+    "portableTileQ": "min(clusterTileQ, pow2ceil(tokenCount), max(1, floor(min(clusterWg, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX) / portableLpq)))",
+    "portableWg": "portableTileQ * portableLpq",
+    "portableKeyCapacity": "max(1, floor(device.limits.maxComputeWorkgroupStorageSize / ((headDim + vHeadDim) * packedStageBytes + 4 * portableWg)))",
+    "portableTileK": "min(clusterTileK, pow2ceil(tokenCount), pow2ceil(portableKeyCapacity + 1) / 2)",
+    "clusterDistributedBytes": "portableTileK * (headDim + vHeadDim) * packedStageBytes + 4 * max(portableWg * portableTileK, 2 * portableWg + portableTileQ * portableTileK)",
+    "clusterDistributedOk": "clusterDistributedBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "clusterUseDistributed": "clusterDistributedOk and (clusterLpq > 4 or vHeadDim != headDim)",
+    "portableTiles": "ceilDiv(tokenCount, portableTileQ if clusterUseDistributed else clusterTileQ)",
+    "clusterPortableOk": "clusterShapeOk and (clusterUseDistributed or clusterNosgOk) and portableWg <= device.limits.maxComputeInvocationsPerWorkgroup and portableWg <= device.limits.maxComputeWorkgroupSizeX and ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
+  },
+  "bindings": {
+    "query": { "arg": "queryT", "name": "query" },
+    "key": { "arg": "keyT", "name": "key" },
+    "value": { "arg": "valueT", "name": "value" },
+    "bias": { "arg": "biasT", "name": "bias" },
+    "tokenOffset": { "arg": "tokenOffsetT", "name": "token_offset" },
+    "cumSeqLen": { "arg": "cumulativeSequenceLengthT", "name": "cumulative_sequence_length" },
+    "attnBias": { "arg": "attentionBiasT", "name": "attn_bias" },
+    "output": { "arg": "outputT", "name": "output" },
+    "query4": { "arg": "queryT", "name": "query", "elementType": "$inputVec4" },
+    "key4": { "arg": "keyT", "name": "key", "elementType": "$inputVec4" },
+    "value4": { "arg": "valueT", "name": "value", "elementType": "$inputVec4" },
+    "output4": { "arg": "outputT", "name": "output", "elementType": "$inputVec4" },
+    "params": {
+      "name": "params",
+      "struct": [
+        { "name": "tokenCount", "type": "u32", "value": "tokenCount" },
+        { "name": "seqPad", "type": "u32", "value": "seqPad" },
+        { "name": "scale", "type": "f32", "value": "attentionScale" }
+      ]
+    },
+    "paramsMask": {
+      "name": "params",
+      "struct": [
+        { "name": "tokenCount", "type": "u32", "value": "tokenCount" },
+        { "name": "seqPad", "type": "u32", "value": "seqPad" },
+        { "name": "scale", "type": "f32", "value": "attentionScale" },
+        { "name": "maskBatchStride", "type": "u32", "value": "maskBatchStride" },
+        { "name": "maskHeadStride", "type": "u32", "value": "maskHeadStride" },
+        { "name": "maskSeqStride", "type": "u32", "value": "maskSeqStride" }
+      ]
+    }
+  },
+  "variants": [
+    {
+      "id": "packed_flash",
+      "priority": 40,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "packed_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "packed_serial",
+      "priority": 0,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "not present.attentionBiasT", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "tokenOffset", "cumSeqLen", "output", "params"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "packed_attn_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "packed_attn_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "packed_attn_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "packedShapeOk", "not present.biasT", "attnBiasOk", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "packed_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "packed_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "packed_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "not present.attentionBiasT", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "bias", "tokenOffset", "cumSeqLen", "output", "params"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "packed_bias_attn_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "packed_bias_attn_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "packed_bias_attn_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "packedShapeOk", "biasOk", "attnBiasOk", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "separate_flash",
+      "priority": 40,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "separate_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "separate_serial",
+      "priority": 0,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "not present.attentionBiasT", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "key", "value", "tokenOffset", "cumSeqLen", "output", "params"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "separate_attn_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "separate_attn_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "separate_attn_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "separateShapeOk", "not present.biasT", "attnBiasOk", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "key", "value", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "separate_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "separate_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "output4", "params"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "separate_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "not present.attentionBiasT", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "key", "value", "bias", "tokenOffset", "cumSeqLen", "output", "params"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "separate_bias_attn_bias_flash",
+      "priority": 40,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "clusterOk", "clusterSubgroupsOk"],
+      "derive": { "useSubgroups": true, "LPQ": "clusterLpq", "TILE_Q": "clusterTileQ", "TILE_K": "clusterTileK" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(clusterTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(clusterTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "requires": { "features": ["subgroups"] },
+      "demoteWhen": ["vHeadDim != headDim and (tokenCount < clusterLpq or (clusterUseDistributed and portableLpq > clusterLpq))"]
+    },
+    {
+      "id": "separate_bias_attn_bias_flash_nosg",
+      "priority": 20,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "clusterPortableOk"],
+      "derive": {
+        "useSubgroups": false,
+        "LPQ": "portableLpq if clusterUseDistributed else clusterLpq",
+        "TILE_Q": "portableTileQ if clusterUseDistributed else clusterTileQ",
+        "TILE_K": "portableTileK if clusterUseDistributed else clusterTileK",
+        "distributedSoftmax": "clusterUseDistributed"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Cluster",
+          "shader": "attn-packed-varlen-cluster.wgsl.jinja",
+          "bindings": ["query4", "key4", "value4", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output4", "paramsMask"],
+          "dispatch": {
+            "x": "min(portableTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(portableTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["vHeadDim != headDim and tokenCount < clusterLpq"]
+    },
+    {
+      "id": "separate_bias_attn_bias_serial",
+      "priority": 0,
+      "when": ["baseContract", "separateShapeOk", "biasOk", "attnBiasOk", "scalarOk"],
+      "derive": { "workgroupSize": "scalarWg" },
+      "passes": [
+        {
+          "id": "main",
+          "name": "PackedMultiHeadAttention.Serial",
+          "shader": "attn-packed-varlen-scalar.wgsl.jinja",
+          "bindings": ["query", "key", "value", "bias", "tokenOffset", "cumSeqLen", "attnBias", "output", "paramsMask"],
+          "dispatch": {
+            "x": "min(tokenCount, DISPATCH_FOLD_WIDTH)",
+            "y": "numHeads",
+            "z": "ceilDiv(tokenCount, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,47 @@

+{
+  "name": "com.microsoft.PackedMultiHeadAttention",
+  "id": "_com_microsoft_packedmultiheadattention_webgpu_96b8ba0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "attn-packed-varlen-cluster.wgsl.jinja": "p9O8Y9FODuK+3sS5XCmt34ubccsQ1rHO8KyYdKFbox8=",
+      "attn-packed-varlen-scalar.wgsl.jinja": "6gm8Q90Uq/Fwm53KBabaxD39SZDE/XPt93lhHeOKEKo=",
+      "bench.json": "QU1tbtEXCd2xke0vrPrixj1rlwWltGZXLr8StqGlVuM=",
+      "manifest.json": "CImYmm6QB0MSplEddT1DpForRg78IeIsHWzBijr+yks=",
+      "test.json": "Y3VjrtppOZI+cmCpPGhsEdWnF8qd981T9VIMTNux648="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "packed_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "packed_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "packed_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "packed_bias_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_bias_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "packed_bias_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "separate_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "separate_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "separate_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"],
+      "separate_bias_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_bias_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja"],
+      "separate_bias_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/test.json ADDED Viewed

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