sync 6fdf6301e2bb

Browse files

Files changed (9) hide show

README.md +132 -0
build/webgpu/attn-packed-varlen-cluster.wgsl.jinja +390 -0
build/webgpu/attn-packed-varlen-scalar.wgsl.jinja +165 -0
build/webgpu/bench.json +0 -0
build/webgpu/gemm-epilogue-tiled-reg.wgsl.jinja +128 -0
build/webgpu/gemm-subgroup-matrix.wgsl.jinja +345 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +61 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.PackedAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Self-attention with a fused Q/K/V input projection 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; there is no mask input. The optional attention bias is indexed in PADDED coordinates. The output is 2-D in packed token order. Float32 and float16 only; bfloat16 is not implemented.
+See the [ONNX Runtime `PackedAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.PackedAttention) for the reference semantics.
+## Inputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `inputT` | `input` | `T` | same as logical dtype | `2` | — | Packed token embeddings with shape `[token_count, input_hidden_size]`. | required |
+| `weightsT` | `weights` | `T` | same as logical dtype | `2` | — | Merged projection weights with shape `[input_hidden_size, q_hidden + k_hidden + v_hidden]`, column-ordered `[Q all heads \| K all heads \| V all heads]`. | required |
+| `biasT` | `bias` | `T` | same as logical dtype | `1` | — | Required projection bias with shape `[q_hidden + k_hidden + v_hidden]`. | required |
+| `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; every Q/K/V hidden size must be divisible by it. |
+| `qkv_hidden_sizes` | — | Optional `[q_hidden, k_hidden, v_hidden]`. Omission splits the bias length in three; `q_hidden` must equal `k_hidden`. |
+| `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.
+- `matrix_direct_f32_flash` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f32_flash` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f32_flash` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs tiled flash attention over the packed token stream: one workgroup covers `CLUSTER_TILE_Q` consecutive tokens for one head and stages `CLUSTER_TILE_K` keys, masking each token to its own sequence. Cluster dot fragments combine through subgroup shuffles.
+- `matrix_direct_f32_flash_nosg` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f32_flash_nosg` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f32_flash_nosg` — Projects Q/K/V with the register-tiled GEMM, then performs packed flash attention. Portable clusters distribute dot fragments and softmax work across lanes. Wider heads use more lanes and fewer queries per workgroup; device limits bound geometry and key staging. Four-lane clusters retain the existing reduction.
+- `matrix_direct_f32_serial` — Projects with dtype-preserving subgroup matrices, direct input reads and double-buffered weights. Float32 uses four accumulation streams. A gated staged pass handles partial rows; sparse grids use smaller staged tiles.
+- `matrix_staged_f32_serial` — Projects with staged 32-row subgroup-matrix tiles and float32 accumulation before packed attention. Small and partial-row grids avoid an extra projection dispatch; reported matrix configurations, subgroup width and resource limits retain portable fallbacks.
+- `f32_serial` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs portable serial-key attention: one workgroup owns one (packed token, head) and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `matrix_direct_f32_attn_bias_flash` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f32_attn_bias_flash` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f32_attn_bias_flash` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs tiled flash attention over the packed token stream: one workgroup covers `CLUSTER_TILE_Q` consecutive tokens for one head and stages `CLUSTER_TILE_K` keys, masking each token to its own sequence. Cluster dot fragments combine through subgroup shuffles.
+- `matrix_direct_f32_attn_bias_flash_nosg` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f32_attn_bias_flash_nosg` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f32_attn_bias_flash_nosg` — Projects Q/K/V with the register-tiled GEMM, then performs packed flash attention. Portable clusters distribute dot fragments and softmax work across lanes. Wider heads use more lanes and fewer queries per workgroup; device limits bound geometry and key staging. Four-lane clusters retain the existing reduction.
+- `matrix_direct_f32_attn_bias_serial` — Projects with dtype-preserving subgroup matrices, direct input reads and double-buffered weights. Float32 uses four accumulation streams. A gated staged pass handles partial rows; sparse grids use smaller staged tiles.
+- `matrix_staged_f32_attn_bias_serial` — Projects with staged 32-row subgroup-matrix tiles and float32 accumulation before packed attention. Small and partial-row grids avoid an extra projection dispatch; reported matrix configurations, subgroup width and resource limits retain portable fallbacks.
+- `f32_attn_bias_serial` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs portable serial-key attention: one workgroup owns one (packed token, head) and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `matrix_direct_f16_flash` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f16_flash` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f16_flash` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs tiled flash attention over the packed token stream: one workgroup covers `CLUSTER_TILE_Q` consecutive tokens for one head and stages `CLUSTER_TILE_K` keys, masking each token to its own sequence. Cluster dot fragments combine through subgroup shuffles.
+- `matrix_direct_f16_flash_nosg` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f16_flash_nosg` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f16_flash_nosg` — Projects Q/K/V with the register-tiled GEMM, then performs packed flash attention. Portable clusters distribute dot fragments and softmax work across lanes. Wider heads use more lanes and fewer queries per workgroup; device limits bound geometry and key staging. Four-lane clusters retain the existing reduction.
+- `matrix_direct_f16_serial` — Projects with dtype-preserving subgroup matrices, direct input reads and double-buffered weights. Float32 uses four accumulation streams. A gated staged pass handles partial rows; sparse grids use smaller staged tiles.
+- `matrix_staged_f16_serial` — Projects with staged 32-row subgroup-matrix tiles and float32 accumulation before packed attention. Small and partial-row grids avoid an extra projection dispatch; reported matrix configurations, subgroup width and resource limits retain portable fallbacks.
+- `f16_serial` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs portable serial-key attention: one workgroup owns one (packed token, head) and reduces each dot through workgroup memory. Head dimensions need no alignment and the value head dimension may differ.
+- `matrix_direct_f16_attn_bias_flash` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f16_attn_bias_flash` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f16_attn_bias_flash` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs tiled flash attention over the packed token stream: one workgroup covers `CLUSTER_TILE_Q` consecutive tokens for one head and stages `CLUSTER_TILE_K` keys, masking each token to its own sequence. Cluster dot fragments combine through subgroup shuffles.
+- `matrix_direct_f16_attn_bias_flash_nosg` — Projects with dtype-preserving matrix tiles and direct input reads; a gated staged pass covers partial rows. Attention uses subgroup or portable tiled reductions according to its own resource and lane budget.
+- `matrix_staged_f16_attn_bias_flash_nosg` — Projects with staged 32-row matrix tiles and float32 accumulation, then runs subgroup or portable tiled attention. Small projection grids avoid a separate tail dispatch; both passes retain independent resource guards.
+- `f16_attn_bias_flash_nosg` — Projects Q/K/V with the register-tiled GEMM, then performs packed flash attention. Portable clusters distribute dot fragments and softmax work across lanes. Wider heads use more lanes and fewer queries per workgroup; device limits bound geometry and key staging. Four-lane clusters retain the existing reduction.
+- `matrix_direct_f16_attn_bias_serial` — Projects with dtype-preserving subgroup matrices, direct input reads and double-buffered weights. Float32 uses four accumulation streams. A gated staged pass handles partial rows; sparse grids use smaller staged tiles.
+- `matrix_staged_f16_attn_bias_serial` — Projects with staged 32-row subgroup-matrix tiles and float32 accumulation before packed attention. Small and partial-row grids avoid an extra projection dispatch; reported matrix configurations, subgroup width and resource limits retain portable fallbacks.
+- `f16_attn_bias_serial` — Fuses the Q/K/V projection into a register-tiled GEMM, then runs portable serial-key attention: one workgroup owns one (packed token, head) 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 `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark 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)
+- [`gemm-epilogue-tiled-reg.wgsl.jinja`](build/webgpu/gemm-epilogue-tiled-reg.wgsl.jinja)
+- [`gemm-subgroup-matrix.wgsl.jinja`](build/webgpu/gemm-subgroup-matrix.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.PackedAttention", { version: 1 });
+const { outputT } = await kernel({
+  inputT: { data: inputTData, shape: [2, 4] },
+  weightsT: { data: weightsTData, shape: [4, 12] },
+  biasT: { data: biasTData, shape: [12] },
+  tokenOffsetT: { data: tokenOffsetTData, shape: [1, 2] },
+  cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [2] },
+}, {
+  attrs: { num_heads: 2 },
+});
+```

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

	@@ -0,0 +1,390 @@

+{% set parallelSoftmax = distributedSoftmax is defined and distributedSoftmax and not useSubgroups %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% set QBUF = "query" %}
+{% set KBUF = "query" %}
+{% set VBUF = "query" %}
+{% 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.
+// Source: one packed `[T, N, 3, H]` buffer; Q, K and V differ only by offset.
+{% 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;
+{% 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;
+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;
+}
+{% 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]);
+  }
+{% else %}
+  var qr{{ c }} = vec4<f32>({{ QBUF }}[qBase4 + {{ c }}u]);
+{% 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]);
+        {{ 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]);
+        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,165 @@

+{{ env.wgsl.resourceDeclarations }}
+{% set QBUF = "query" %}
+{% set KBUF = "query" %}
+{% set VBUF = "query" %}
+// 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.
+// Source: one packed `[T, N, 3, H]` buffer; Q, K and V differ only by offset.
+{% 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;
+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]);
+      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]);
+      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/gemm-epilogue-tiled-reg.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,128 @@

+{% set OUT = "f16" if (usesF16 is defined and usesF16) else "f32" %}
+{{ env.wgsl.resourceDeclarations }}
+// Shape-driven, register-blocked f32 GEMM for all four transpose combinations.
+// Runtime M/N/K share one compiled pipeline for each tile configuration.
+// Register micro-tiles reuse staged K slices; transpose modes change load strides and
+// keep the stored-contiguous axis coalesced. M/N/K tails are bounds-checked
+// unless an authored dimension proves that its tile has no tail.
+// Both tiles are indexed by their own output axis and group four K values per
+// vector word, so the micro-tile accumulates through dot() and one step reads
+// TM + TN words instead of 4 * (TM + TN) scalars.
+{% set epiOpen = "" %}
+{% set epiClose = "" %}
+{% set gemmWgX = ((gemmTileCols / gemmMicroCols) | int) %}
+{% set gemmWgY = ((gemmTileRows / gemmMicroRows) | int) %}
+{% set gemmTileK = 16 %}
+const BK: u32 = {{ gemmTileK }}u;
+const BM: u32 = {{ gemmTileRows }}u;
+const BN: u32 = {{ gemmTileCols }}u;
+const TM: u32 = {{ gemmMicroRows }}u;   // per-thread micro-tile rows
+const TN: u32 = {{ gemmMicroCols }}u;   // per-thread micro-tile cols
+// Stage tiles at the operands' storage width. For f16 inputs this is
+// bit-identical to widening on shared load, keeps f32 FMA registers, and halves
+// workgroup storage.
+{% set tileT = OUT %}
+const K_VECS: u32 = BK / 4u;
+var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // logical A'[m][k/4]
+var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // logical B'[n][k/4]
+@compute @workgroup_size({{ gemmWgX }}, {{ gemmWgY }}, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let li = lid.y * {{ gemmWgX }}u + lid.x;
+  let M = params.M;
+  let N = params.N;
+  let K = params.K;
+  let n_tiles = (N + BN - 1u) / BN;
+  let m_tiles = (M + BM - 1u) / BM;
+  for (var nt: u32 = wg.x; nt < n_tiles; nt = nt + nwg.x) {
+    let nBase = nt * BN;
+    for (var mt: u32 = wg.y; mt < m_tiles; mt = mt + nwg.y) {
+      let mBase = mt * BM;
+      var acc: array<f32, TM * TN>;
+      for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
+      let numTiles = (K + BK - 1u) / BK;
+      for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+        let kBase = kt * BK;
+        // Cooperative load: one vector word per lane per pass. The lane grid is
+        // laid out so the axis the operand stores contiguously stays coalesced.
+        for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + {{ gemmWgX * gemmWgY }}u) {
+          // A stored [M, K]: logical A[m,k] = a[m*K + k]. Coalesce on k (contiguous).
+          let ar = idx / K_VECS;
+          let ac4 = idx % K_VECS;
+          let am = mBase + ar;
+          let ak = kBase + ac4 * 4u;
+          var aWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+          if (am < M) {
+{% for component in range(4) %}
+            if (ak + {{ component }}u < K) {
+              aWord[{{ component }}u] = {{ tileT }}(a[am * K + ak + {{ component }}u]);
+            }
+{% endfor %}
+          }
+          tileA[ar][ac4] = aWord;
+        }
+        for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + {{ gemmWgX * gemmWgY }}u) {
+          // B stored [K, N]: logical B[k,n] = b[k*N + n]. Coalesce on n (contiguous).
+          let bc = idx % BN;
+          let br4 = idx / BN;
+          let bn = nBase + bc;
+          let bk = kBase + br4 * 4u;
+          var bWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+          if (bn < N) {
+{% for component in range(4) %}
+            if (bk + {{ component }}u < K) {
+              bWord[{{ component }}u] = {{ tileT }}(b[(bk + {{ component }}u) * N + bn]);
+            }
+{% endfor %}
+          }
+          tileB[bc][br4] = bWord;
+        }
+        workgroupBarrier();
+{% set regAccumulator = "acc" %}
+{% set regT = "f32" %}{% set regCast = 1 if tileT != "f32" else 0 %}
+{% filter indent(8, true) %}
+{% set regAccumulator = "acc" %}
+let aRow = lid.y * TM;
+let bCol = lid.x * TN;
+for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+  var av: array<vec4<{{ regT }}>, TM>;
+  var bv: array<vec4<{{ regT }}>, TN>;
+  for (var i: u32 = 0u; i < TM; i = i + 1u) { av[i] = {% if regCast %}vec4<{{ regT }}>(tileA[aRow + i][kv]){% else %}tileA[aRow + i][kv]{% endif %}; }
+  for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = {% if regCast %}vec4<{{ regT }}>(tileB[bCol + j][kv]){% else %}tileB[bCol + j][kv]{% endif %}; }
+  for (var i: u32 = 0u; i < TM; i = i + 1u) {
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + dot(av[i], bv[j]);
+    }
+  }
+}
+{% endfilter %}
+        workgroupBarrier();
+      }
+      let m0 = mBase + lid.y * TM;
+      let n0 = nBase + lid.x * TN;
+      for (var ti: u32 = 0u; ti < TM; ti = ti + 1u) {
+        let m = m0 + ti;
+        if (m >= M) { continue; }
+        for (var tj: u32 = 0u; tj < TN; tj = tj + 1u) {
+          let n = n0 + tj;
+          if (n < N) {
+            // alpha == 0 nulls A*B exactly and avoids Inf*0 becoming NaN.
+            let product = select(acc[ti * TN + tj] * params.alpha, 0.0, params.alpha == 0.0);
+            let biasValue = f32(bias[n]);
+            let biasTerm = select(biasValue * params.beta, 0.0, params.beta == 0.0);
+            y[m * N + n] = {{ OUT }}({{ epiOpen }}product + biasTerm{{ epiClose }});
+          }
+        }
+      }
+    }
+  }
+}

build/webgpu/gemm-subgroup-matrix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,345 @@

+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+{% set operandScalar = mmaScalar %}
+{% set accScalar = "f32" %}
+{% set OUT = "y" %}
+{% set OUT_SCALAR = outScalar if outScalar is defined else yScalar %}
+{% set tileMValue = tileM if tileM is defined else 32 %}
+{% set tileNValue = tileN if tileN is defined else 64 %}
+{% set tileKValue = 32 %}
+{% set epiOpen = "" %}
+{% set epiClose = "" %}
+{% set useDirectStore = false %}
+{% set useDirectAStorage = directAStorage is defined and directAStorage %}
+{% set useDoubleBufferedB = doubleBufferedB is defined and doubleBufferedB %}
+{% set useBalancedAccumulation = balancedMatrixAccumulation is defined and balancedMatrixAccumulation and useDoubleBufferedB and not useDirectStore and operandScalar == "f32" %}
+{% set columnTileBaseValue = 0 %}
+{% set rowTileBaseValue = rowTileBase if rowTileBase is defined else 0 %}
+// Four subgroups cover 32xN as a 2x2 grid, or taller tiles as four row bands.
+{% if tileMValue > 32 %}
+{% set subtileCols = tileNValue %}
+{% set subtileRows = (tileMValue / 4)|int %}
+{% else %}
+{% set subtileCols = (tileNValue / 2)|int %}
+{% set subtileRows = 16 %}
+{% endif %}
+{% set aTilesPerSubgroup = (subtileRows / 8)|int %}
+{% set bTilesPerSubgroup = (subtileCols / 8)|int %}
+{% set subgroupCount = 4 %}
+{% set scratchBanks = subgroupCount * aTilesPerSubgroup * bTilesPerSubgroup %}
+{% macro matrixAccumulators() %}
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  var matC{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% if useBalancedAccumulation %}
+  var matD{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matE{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+  var matF{{ m }}{{ n }}: subgroup_matrix_result<{{ accScalar }}, 8, 8>;
+{% endif %}
+{% endfor %}
+{% endfor %}
+{% endmacro %}
+{% macro storeMatrixAccumulators() %}
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  subgroupMatrixStore<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matC{{ m }}{{ n }}, 8u);
+{% if useBalancedAccumulation %}
+  subgroupMatrixStore<row_major>(&tile_B, (bank + {{ m * bTilesPerSubgroup + n }}u) * 64u, matD{{ m }}{{ n }}, 8u);
+{% endif %}
+{% endfor %}
+{% endfor %}
+{% endmacro %}
+const IN_F:       u32 = {{ inFeatures }}u;
+const OUT_F:      u32 = {{ outFeatures }}u;
+{% set kLoopEndValue = "IN_F" %}
+const TILE_COLS:  u32 = {{ tileNValue }}u;
+const TILE_ROWS:  u32 = {{ tileMValue }}u;
+const TILE_K:     u32 = {{ tileKValue }}u;
+const SUB_COLS:   u32 = {{ subtileCols }}u;
+const SUB_ROWS:   u32 = {{ subtileRows }}u;
+const COLUMN_TILE_BASE: u32 = {{ columnTileBaseValue }}u;
+const ROW_TILE_BASE: u32 = {{ rowTileBaseValue }}u;
+{% if not useDirectAStorage %}
+var<workgroup> tile_A: array<{{ operandScalar }}, {{ tileMValue }} * {{ tileKValue }}>;
+{% endif %}
+var<workgroup> tile_B: array<{{ operandScalar }}, {{ (2 if useDoubleBufferedB else 1) * tileNValue }} * {{ tileKValue }}>;
+// Distinct readback banks for every result matrix, so one barrier can publish
+// the full subtile without write-after-read reuse.
+var<workgroup> scratch: array<array<{{ accScalar }}, 64>, {{ scratchBanks }}>;
+{% if not useDirectAStorage %}
+fn loadSHMA(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
+  // Clamp the partial-M tail to the last valid row (M-1) instead of zero-filling
+  // it. Some subgroup-matrix implementations let a zero-padded left tile corrupt
+  // the adjacent valid row. Replicating row M-1 avoids that boundary; the
+  // replicated rows are discarded by storeOutput's row_limit guard, so every valid
+  // output row stays bit-identical. M >= 1 is guaranteed by the when-clause args.M > 0.
+  let col: u32 = c_idx * 8u;
+  for (var row_offset: u32 = 0u; row_offset < TILE_ROWS; row_offset += {{ (1024 / tileKValue)|int }}u) {
+  let r: u32 = row + row_offset;
+  let a_global: u32 = min(tile_base + r, params.M - 1u);
+  for (var col_offset: u32 = 0u; col_offset < 8u; col_offset++) {
+    let k: u32 = k_idx + col + col_offset;
+    tile_A[r * TILE_K + col + col_offset] =
+      x[a_global * IN_F + k];
+  }
+  }
+}
+{% endif %}
+fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32, bank_offset: u32) {
+  // Generic transB=0 stores B as [K,N]. Stage that native orientation and load
+  // the subgroup right operand without a transpose.
+  let local_idx = row * {{ (tileKValue / 16)|int }}u + c_idx;
+  let vectors_per_k = TILE_COLS / 4u;
+  let total_vectors = TILE_K * vectors_per_k;
+  for (var vector_idx = local_idx; vector_idx < total_vectors; vector_idx += 128u) {
+    let k_local = vector_idx / vectors_per_k;
+    let n4 = (vector_idx - k_local * vectors_per_k) * 4u;
+    for (var component = 0u; component < 4u; component++) {
+      let n_local = n4 + component;
+      let global_n = tile_base + n_local;
+      let global_k = k_idx + k_local;
+      tile_B[bank_offset + k_local * TILE_COLS + n_local] = w[global_k * OUT_F + global_n];
+    }
+  }
+}
+{% set BIAS_1 = " + f32(bias[col_base + col])" %}
+{% set BIAS_2 = " + f32(bias[col_base + col2])" %}
+{% set ACC_1 = "(prior.x + (scratch[src_slot][row * 8u + col] + tile_B[src_slot * 64u + row * 8u + col]))" if useBalancedAccumulation else "scratch[src_slot][row * 8u + col]" %}
+{% set ACC_2 = "(prior.y + (scratch[src_slot][row * 8u + col2] + tile_B[src_slot * 64u + row * 8u + col2]))" if useBalancedAccumulation else "scratch[src_slot][row * 8u + col2]" %}
+fn storeOutput(offset: u32, col_base: u32, row: u32, col: u32, src_slot: u32, row_limit: i32{% if useBalancedAccumulation %}, prior: vec2<f32>{% endif %}) {
+  if (row_limit > 0 && row < u32(row_limit)) {
+    let col2: u32 = col + 1u;
+    {{ OUT }}[offset + row * OUT_F + col]  = {{ OUT_SCALAR }}({{ epiOpen }}{{ ACC_1 }}{{ BIAS_1 }}{{ epiClose }});
+    {{ OUT }}[offset + row * OUT_F + col2] = {{ OUT_SCALAR }}({{ epiOpen }}{{ ACC_2 }}{{ BIAS_2 }}{{ epiClose }});
+  }
+}
+@compute @workgroup_size(128, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_invocation_id) sg_id: u32,
+  @builtin(subgroup_size) sg_size: u32
+) {
+  // 2D-folded M-tile (row) base: workgroup_id.z carries the high bits when the
+  // row tile exceeds the per-dimension dispatch limit.
+  // Reduces to workgroup_id.y when nwg.z == 1; the row_limit guard in storeOutput
+  // (i32(M) - i32(a_global_base + ...)) drops the over-dispatched tail.
+  let a_global_base: u32 = (ROW_TILE_BASE + workgroup_id.y + workgroup_id.z * nwg.y) * TILE_ROWS;
+  let w_global_base: u32 = (COLUMN_TILE_BASE + workgroup_id.x) * TILE_COLS;
+  let subtile_id: u32 = local_idx / sg_size;
+{% if tileMValue > 32 %}
+  let subtile_idx: u32 = 0u;
+  let subtile_idy: u32 = subtile_id;
+{% else %}
+  let subtile_idx: u32 = subtile_id / 2u;
+  let subtile_idy: u32 = subtile_id % 2u;
+{% endif %}
+  let base_A: u32 = subtile_idy * SUB_ROWS;
+  let base_B: u32 = subtile_idx * SUB_COLS;
+{{ matrixAccumulators() }}
+{% if useDoubleBufferedB %}
+  // Large aligned f32 tiles do not benefit from staging A: each row band is
+  // consumed by exactly one subgroup. Load those matrices directly from x and
+  // ping-pong only B, which is reused by all four subgroups. The barrier that
+  // publishes the next B bank also proves the previous bank is no longer read,
+  // reducing the steady-state schedule from two barriers per K tile to one.
+  var b_bank: u32 = 0u;
+  loadSHMB(w_global_base, 0u, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, 0u);
+  workgroupBarrier();
+  for (var kidx: u32 = 0u; kidx < {{ kLoopEndValue }}; kidx += TILE_K) {
+    let b_bank_offset = b_bank * TILE_COLS * TILE_K;
+    for (var step: u32 = 0u; step < TILE_K; step += {{ 32 if useBalancedAccumulation else 8 }}u) {
+{% if useBalancedAccumulation %}
+{% for stream in range(4) %}
+    {
+{% set matrixAccumulator = ["matC", "matD", "matE", "matF"][stream] %}
+{% set matrixStepOffset = stream * 8 %}
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set directA = directAStorage is defined and directAStorage %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
+{% else %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
+{% endfor %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
+{% for n in range(bTilesPerSubgroup) %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
+{% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+    }
+{% endfor %}
+{% else %}
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set directA = directAStorage is defined and directAStorage %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
+{% else %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
+{% endfor %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
+{% for n in range(bTilesPerSubgroup) %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
+{% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+{% endif %}
+    }
+    let next_kidx = kidx + TILE_K;
+    if (next_kidx < {{ kLoopEndValue }}) {
+      let next_bank = 1u - b_bank;
+      loadSHMB(w_global_base, next_kidx, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, next_bank * TILE_COLS * TILE_K);
+      workgroupBarrier();
+      b_bank = next_bank;
+    }
+  }
+{% else %}
+  for (var kidx: u32 = 0u; kidx < {{ kLoopEndValue }}; kidx += TILE_K) {
+{% if not useDirectAStorage %}
+    loadSHMA(a_global_base, kidx, local_idx / {{ (tileKValue / 8)|int }}u, local_idx % {{ (tileKValue / 8)|int }}u);
+{% endif %}
+    loadSHMB(w_global_base, kidx, local_idx / {{ (tileKValue / 16)|int }}u, local_idx % {{ (tileKValue / 16)|int }}u, 0u);
+    workgroupBarrier();
+    for (var step: u32 = 0u; step < TILE_K; step += 8u) {
+{% set dynamicATiles = aTilesPerSubgroup if aTilesPerSubgroup is defined else 2 %}
+{% set matrixAccumulator = matrixAccumulator | default("matC") %}
+{% set matrixStepOffset = matrixStepOffset | default(0) %}
+{% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
+{% set directA = directAStorage is defined and directAStorage %}
+{% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
+{% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
+      // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
+      // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
+{% else %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
+{% endif %}
+{% for m in range(dynamicATiles) %}
+{% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
+{% else %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
+{% endif %}
+{% endfor %}
+      let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
+{% for n in range(bTilesPerSubgroup) %}
+      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
+{% endfor %}
+{% for m in range(dynamicATiles) %}
+{% for n in range(bTilesPerSubgroup) %}
+      {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% endif %}
+{% if useBalancedAccumulation %}
+  workgroupBarrier();
+{% endif %}
+  let matrix_c_offset: u32 = (a_global_base + base_A) * OUT_F + w_global_base + base_B;
+  let row: u32 = sg_id / 4u;
+  let col: u32 = (sg_id % 4u) * 2u;
+  // Stage every result into a distinct bank before one barrier. Tail/bias paths
+  // retain this guarded scalar epilogue: subgroupMatrixStore scatters values
+  // across lanes, so its cross-lane readback must be published before partial-M
+  // guards diverge. The epilogue also handles output conversion and bias.
+  let bank: u32 = subtile_id * {{ aTilesPerSubgroup * bTilesPerSubgroup }}u;
+{{ storeMatrixAccumulators() }}
+  workgroupBarrier();
+{% if useBalancedAccumulation %}
+  var pairTotals: array<vec2<f32>, {{ aTilesPerSubgroup * bTilesPerSubgroup }}>;
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  pairTotals[{{ m * bTilesPerSubgroup + n }}] = vec2<f32>(
+    scratch[bank + {{ m * bTilesPerSubgroup + n }}u][row * 8u + col] + tile_B[(bank + {{ m * bTilesPerSubgroup + n }}u) * 64u + row * 8u + col],
+    scratch[bank + {{ m * bTilesPerSubgroup + n }}u][row * 8u + col + 1u] + tile_B[(bank + {{ m * bTilesPerSubgroup + n }}u) * 64u + row * 8u + col + 1u]);
+{% endfor %}
+{% endfor %}
+  workgroupBarrier();
+{% for m in range(aTilesPerSubgroup) %}
+{% for n in range(bTilesPerSubgroup) %}
+  subgroupMatrixStore<row_major>(&scratch[bank + {{ m * bTilesPerSubgroup + n }}u], 0u, matE{{ m }}{{ n }}, 8u);
+  subgroupMatrixStore<row_major>(&tile_B, (bank + {{ m * bTilesPerSubgroup + n }}u) * 64u, matF{{ m }}{{ n }}, 8u);
+{% endfor %}
+{% endfor %}
+  workgroupBarrier();
+{% endif %}
+{% for m in range(aTilesPerSubgroup) %}
+  let row_limit_{{ m }}: i32 = i32(params.M) - i32(a_global_base + base_A + {{ m * 8 }}u);
+{% for n in range(bTilesPerSubgroup) %}
+  storeOutput(
+    matrix_c_offset + {{ m * 8 }}u * OUT_F + {{ n * 8 }}u,
+    w_global_base + base_B + {{ n * 8 }}u,
+    row,
+    col,
+    bank + {{ m * bTilesPerSubgroup + n }}u,
+    row_limit_{{ m }}{{ ", pairTotals[" ~ (m * bTilesPerSubgroup + n) ~ "]" if useBalancedAccumulation else "" }}
+  );
+{% endfor %}
+{% endfor %}
+}

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,61 @@

+{
+  "name": "com.microsoft.PackedAttention",
+  "id": "_com_microsoft_packedattention_webgpu_2e383dd",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "attn-packed-varlen-cluster.wgsl.jinja": "M80JSlkb8TDwWae5JchL6XbpV286kAlQj5Vzibog8bQ=",
+      "attn-packed-varlen-scalar.wgsl.jinja": "tHu3bLC21Ptmb1VpMTt4ztk0sYa/zMjAqMGohlJGGgk=",
+      "bench.json": "8ZMv62kt75ShLjFueX/PRLS3gqcBHFs4BNOVE8RvzP8=",
+      "gemm-epilogue-tiled-reg.wgsl.jinja": "qQqLZ53/gvjjKWGMEydTYCqjBIX0DVYF1kGPt/tb0F4=",
+      "gemm-subgroup-matrix.wgsl.jinja": "nGflGDEU4lkkKKhH8SgascRjubigTw8iLGgHMUDqLV4=",
+      "manifest.json": "0UPOkWWzNh4GybUb6c4Nf2UNPNl6Lqb0t/mM10I6KTQ=",
+      "test.json": "Q/fxohWe7N1k2sg8Av4jtPf7tSvT3AEp1MERMj6KMRI="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "matrix_direct_f32_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f32_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f32_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f32_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f32_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f32_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f32_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f32_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_attn_bias_flash": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_attn_bias_flash_nosg": ["attn-packed-varlen-cluster.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "matrix_direct_f16_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "matrix_staged_f16_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-subgroup-matrix.wgsl.jinja"],
+      "f16_attn_bias_serial": ["attn-packed-varlen-scalar.wgsl.jinja", "gemm-epilogue-tiled-reg.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/test.json ADDED Viewed

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