sync 2e7068faf55e

Browse files

Files changed (12) hide show

README.md +91 -0
build/webgpu/bench.json +323 -0
build/webgpu/chunk-out.wgsl.jinja +200 -0
build/webgpu/chunk-prep.wgsl.jinja +47 -0
build/webgpu/chunk-scan.wgsl.jinja +216 -0
build/webgpu/chunk-ut.wgsl.jinja +234 -0
build/webgpu/linear-attention.scalar.wgsl.jinja +332 -0
build/webgpu/linear-attention.serial.wgsl.jinja +155 -0
build/webgpu/linear-attention.vec4.wgsl.jinja +349 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +24 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.LinearAttention
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Recurrent linear attention for packed `[B, T, H*D]` decode and prefill. It supports all four update rules, standard and inverse GQA, shared-key heads, and rollback states through `state_window`. Activations and state may independently use float16 or float32; bfloat16 is not implemented. `past_state` is optional for every update rule and defaults to zeros.
+See the [ONNX Runtime `LinearAttention` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.LinearAttention) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `query` | `queryT` | `T` | `3` | — | Query vectors with 3D packed shape `(B, T, H_q * d_k)`; heads are packed into the last dimension. | required |
+| `key` | `keyT` | `T` | `3` | — | Key vectors with 3D packed shape `(B, T, H_k * d_k)`, where positive `H_k` divides `H_kv`; `H_k < H_kv` shares each key head across multiple KV-state heads. Keys should be L2-normalized for `delta`/`gated_delta` modes. | required |
+| `value` | `valueT` | `T` | `3` | — | Value vectors with 3D packed shape `(B, T, H_kv * d_v)`. | required |
+| `past_state` | `pastStateT` | `S` | derived | derived; see description | Recurrent state from the previous step with shape `(B, H_kv, d_k, d_v)`, or `(W, B, H_kv, d_k, d_v)` when `state_window = W > 0`; defaults to zeros if absent. | optional |
+| `decay` | `decayT` | `T` | `3` | — | Exponential decay gate in log-space with shape `(B, T, H_kv * d_k)` or `(B, T, H_kv)`; required for `gated` and `gated_delta` modes. | optional |
+| `beta` | `betaT` | `T` | `3` | — | Update rate (sigmoid output) with shape `(B, T, H_kv)` or `(B, T, 1)`; required for `delta` and `gated_delta` modes. | optional |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `outputT` | `T` | `3` | derived; see description | Attention output with 3D packed shape `(B, T, max(H_q, H_kv) * d_v)`. | required |
+| `present_state` | `presentStateT` | `S` | derived | derived; see description | Updated recurrent state with shape `(B, H_kv, d_k, d_v)`, or `(W, B, H_kv, d_k, d_v)` when `state_window = W > 0`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `chunk_size` | `64` | Accepted for schema compatibility; does not affect the result. |
+| `scale` | `0` | Scale applied to query-key products. Zero selects `1 / sqrt(d_k)`. |
+| `state_window` | `0` | Number of recent recurrent states retained in `present_state`, in the supported range 0 to 8; zero returns only the current state. |
+| `update_rule` | `"gated_delta"` | Recurrent update rule: `linear`, `gated`, `delta`, or `gated_delta`. |
+| `kv_num_heads` | — | Number of key/value heads. |
+| `q_num_heads` | — | Number of query heads. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `S` | `float32`, `float16` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`chunk-out.wgsl.jinja`](build/webgpu/chunk-out.wgsl.jinja)
+- [`chunk-prep.wgsl.jinja`](build/webgpu/chunk-prep.wgsl.jinja)
+- [`chunk-scan.wgsl.jinja`](build/webgpu/chunk-scan.wgsl.jinja)
+- [`chunk-ut.wgsl.jinja`](build/webgpu/chunk-ut.wgsl.jinja)
+- [`linear-attention.scalar.wgsl.jinja`](build/webgpu/linear-attention.scalar.wgsl.jinja)
+- [`linear-attention.serial.wgsl.jinja`](build/webgpu/linear-attention.serial.wgsl.jinja)
+- [`linear-attention.vec4.wgsl.jinja`](build/webgpu/linear-attention.vec4.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.LinearAttention", { version: 1 });
+const { outputT, presentStateT } = await kernel({
+  queryT: { data: queryTData, shape: [1, 3, 8] },
+  keyT: { data: keyTData, shape: [1, 3, 4] },
+  valueT: { data: valueTData, shape: [1, 3, 4] },
+  pastStateT: { data: pastStateTData, shape: [1, 1, 4, 4] },
+  decayT: { data: decayTData, shape: [1, 3, 1] },
+  betaT: { data: betaTData, shape: [1, 3, 1] },
+}, {
+  attrs: { q_num_heads: 2, kv_num_heads: 1 },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,323 @@

+{
+  "op": "com.microsoft.LinearAttention",
+  "tunableSpace": {
+    "dvGroups": [2, 4, 8],
+    "tileV": [4, 8, 16],
+    "gatedTileV": [2, 4, 8],
+    "chunkSize": [16, 32],
+    "chunkTileV": [16, 32]
+  },
+  "cases": [
+    {
+      "name": "linear-attention-f32-zero-32x4x16x16",
+      "preset": "smoke",
+      "vars": { "batch": 1, "seq": 32, "qHeads": 4, "kvHeads": 2, "dk": 16, "dv": 16 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "linear", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [1, 32, 64], "dtype": "float32", "dist": "normal", "seed": 201, "scale": 0.2 },
+        "keyT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 202, "scale": 0.2 },
+        "valueT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 203, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 32, 64], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 2, 16, 16], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "2 * args.batch * args.seq * (args.qHeads + args.kvHeads) * args.dk * args.dv" }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-f32-state-32x4x16x16",
+      "preset": "smoke",
+      "vars": { "batch": 1, "seq": 32, "qHeads": 4, "kvHeads": 2, "dk": 16, "dv": 16 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "gated_delta", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [1, 32, 64], "dtype": "float32", "dist": "normal", "seed": 202, "scale": 0.2 },
+        "keyT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 203, "scale": 0.2 },
+        "valueT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 204, "scale": 0.2 },
+        "pastStateT": { "shape": [1, 2, 16, 16], "dtype": "float32", "dist": "normal", "seed": 205, "scale": 0.1 },
+        "decayT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 206, "scale": 0.1 },
+        "betaT": { "shape": [1, 32, 2], "dtype": "float32", "dist": "normal", "seed": 207, "scale": 0.1 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 32, 64], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 2, 16, 16], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-f32-linear-state-32x4x16x16",
+      "preset": "smoke",
+      "vars": { "batch": 1, "seq": 32, "qHeads": 4, "kvHeads": 2, "dk": 16, "dv": 16 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "linear", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [1, 32, 64], "dtype": "float32", "dist": "normal", "seed": 208, "scale": 0.2 },
+        "keyT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 209, "scale": 0.2 },
+        "valueT": { "shape": [1, 32, 32], "dtype": "float32", "dist": "normal", "seed": 210, "scale": 0.2 },
+        "pastStateT": { "shape": [1, 2, 16, 16], "dtype": "float32", "dist": "normal", "seed": 211, "scale": 0.1 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 32, 64], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 2, 16, 16], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "2 * args.batch * args.seq * (args.qHeads + args.kvHeads) * args.dk * args.dv" }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-linear-state-scalar-f16-seq1536-pathology",
+      "preset": "stress",
+      "provenance": {
+        "source": "authored for branch coverage",
+        "notes": "Long-sequence supplied-state case at head_dim_k 16. It exercises the recurrent small-dk route's serial token recurrence and distinguishes it from the chunked prefill decomposition."
+      },
+      "vars": { "batch": 4, "seq": 1536, "qHeads": 4, "kvHeads": 2, "dk": 16, "dv": 16 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "linear", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [4, 1536, 64], "dtype": "float16", "dist": "normal", "seed": 314, "scale": 0.2 },
+        "keyT": { "shape": [4, 1536, 32], "dtype": "float16", "dist": "normal", "seed": 315, "scale": 0.2 },
+        "valueT": { "shape": [4, 1536, 32], "dtype": "float16", "dist": "normal", "seed": 316, "scale": 0.2 },
+        "pastStateT": { "shape": [4, 2, 16, 16], "dtype": "float16", "dist": "normal", "seed": 317, "scale": 0.1 }
+      },
+      "outputs": {
+        "outputT": { "shape": [4, 1536, 64], "dtype": "float16" },
+        "presentStateT": { "shape": [4, 2, 16, 16], "dtype": "float16" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "2 * args.batch * args.seq * (args.qHeads + args.kvHeads) * args.dk * args.dv" }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-gated_delta-scalar-headdimk6-seq1536-stress",
+      "preset": "stress",
+      "vars": { "batch": 8, "seq": 1536, "qHeads": 4, "kvHeads": 4, "dk": 6, "dv": 12 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 4, "update_rule": "gated_delta", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [8, 1536, 24], "dtype": "float32", "dist": "normal", "seed": 301, "scale": 0.2 },
+        "keyT": { "shape": [8, 1536, 24], "dtype": "float32", "dist": "normal", "seed": 302, "scale": 0.2 },
+        "valueT": { "shape": [8, 1536, 48], "dtype": "float32", "dist": "normal", "seed": 303, "scale": 0.2 },
+        "pastStateT": { "shape": [8, 4, 6, 12], "dtype": "float32", "dist": "normal", "seed": 304, "scale": 0.1 },
+        "decayT": { "shape": [8, 1536, 4], "dtype": "float32", "dist": "normal", "seed": 305, "scale": 0.1 },
+        "betaT": { "shape": [8, 1536, 4], "dtype": "float32", "dist": "normal", "seed": 306, "scale": 0.1 }
+      },
+      "outputs": {
+        "outputT": { "shape": [8, 1536, 48], "dtype": "float32" },
+        "presentStateT": { "shape": [8, 4, 6, 12], "dtype": "float32" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-linear-scalar-f16-seq1536-stress",
+      "preset": "stress",
+      "vars": { "batch": 4, "seq": 1536, "qHeads": 4, "kvHeads": 2, "dk": 16, "dv": 16 },
+      "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "linear", "scale": 0.25 },
+      "inputs": {
+        "queryT": { "shape": [4, 1536, 64], "dtype": "float16", "dist": "normal", "seed": 311, "scale": 0.2 },
+        "keyT": { "shape": [4, 1536, 32], "dtype": "float16", "dist": "normal", "seed": 312, "scale": 0.2 },
+        "valueT": { "shape": [4, 1536, 32], "dtype": "float16", "dist": "normal", "seed": 313, "scale": 0.2 }
+      },
+      "outputs": {
+        "outputT": { "shape": [4, 1536, 64], "dtype": "float16" },
+        "presentStateT": { "shape": [4, 2, 16, 16], "dtype": "float16" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          { "type": "gflops", "value": "2 * args.batch * args.seq * (args.qHeads + args.kvHeads) * args.dk * args.dv" }
+        ]
+      },
+      "provenance": {
+        "source": "authored for branch coverage",
+        "notes": "Zero-state sibling of the supplied-state long-sequence case, covering the same branch and shape without an entry state."
+      }
+    },
+    {
+      "name": "linear-attention-gated-delta-f32-bonsai-m16-h48-kv16-dk128-dv128",
+      "preset": "stress",
+      "vars": { "batch": 1, "seq": 16, "qHeads": 48, "kvHeads": 16, "dk": 128, "dv": 128 },
+      "attrs": { "q_num_heads": 48, "kv_num_heads": 16, "update_rule": "gated_delta", "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": { "shape": [1, 16, 6144], "dtype": "float32", "dist": "normal", "seed": 320, "scale": 0.05 },
+        "keyT": { "shape": [1, 16, 2048], "dtype": "float32", "dist": "normal", "seed": 321, "scale": 0.05 },
+        "valueT": { "shape": [1, 16, 2048], "dtype": "float32", "dist": "normal", "seed": 322, "scale": 0.05 },
+        "pastStateT": { "shape": [1, 16, 128, 128], "dtype": "float32", "dist": "normal", "seed": 323, "scale": 0.02 },
+        "decayT": { "shape": [1, 16, 16], "dtype": "float32", "dist": "normal", "seed": 324, "scale": 0.08 },
+        "betaT": { "shape": [1, 16, 16], "dtype": "float32", "dist": "normal", "seed": 325, "scale": 0.08 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 16, 6144], "dtype": "float32", "dist": "empty" },
+        "presentStateT": { "shape": [1, 16, 128, 128], "dtype": "float32", "dist": "empty" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-gated-delta-f16-bonsai-m16-h48-kv16-dk128-dv128",
+      "preset": "stress",
+      "vars": { "batch": 1, "seq": 16, "qHeads": 48, "kvHeads": 16, "dk": 128, "dv": 128 },
+      "attrs": { "q_num_heads": 48, "kv_num_heads": 16, "update_rule": "gated_delta", "scale": 0.08838834764831845 },
+      "inputs": {
+        "queryT": { "shape": [1, 16, 6144], "dtype": "float16", "dist": "normal", "seed": 326, "scale": 0.05 },
+        "keyT": { "shape": [1, 16, 2048], "dtype": "float16", "dist": "normal", "seed": 327, "scale": 0.05 },
+        "valueT": { "shape": [1, 16, 2048], "dtype": "float16", "dist": "normal", "seed": 328, "scale": 0.05 },
+        "pastStateT": { "shape": [1, 16, 128, 128], "dtype": "float16", "dist": "normal", "seed": 329, "scale": 0.02 },
+        "decayT": { "shape": [1, 16, 16], "dtype": "float16", "dist": "normal", "seed": 330, "scale": 0.08 },
+        "betaT": { "shape": [1, 16, 16], "dtype": "float16", "dist": "normal", "seed": 331, "scale": 0.08 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 16, 6144], "dtype": "float16", "dist": "empty" },
+        "presentStateT": { "shape": [1, 16, 128, 128], "dtype": "float16", "dist": "empty" }
+      },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-qwen3next-decode-s1",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-Next class defaults (linear_num_value_heads 32, linear_num_key_heads 16, linear_key_head_dim 128, linear_value_head_dim 128) at a decode step, where the recurrence carries the whole cost."
+      },
+      "vars": { "batch": 1, "seq": 1, "qHeads": 32, "kvHeads": 16, "dk": 128, "dv": 128 },
+      "attrs": {
+        "q_num_heads": 32,
+        "kv_num_heads": 16,
+        "update_rule": "gated_delta",
+        "scale": 0.08838834764831843,
+        "chunk_size": 64
+      },
+      "inputs": {
+        "queryT": { "shape": [1, 1, 4096], "dtype": "float32", "dist": "normal", "seed": 8100, "scale": 0.3 },
+        "keyT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 8101, "scale": 0.3 },
+        "valueT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "normal", "seed": 8102, "scale": 0.3 },
+        "pastStateT": { "shape": [1, 16, 128, 128], "dtype": "float32", "dist": "normal", "seed": 8103, "scale": 0.1 },
+        "decayT": { "shape": [1, 1, 2048], "dtype": "float32", "dist": "uniform", "seed": 8104, "min": 0.9, "max": 1 },
+        "betaT": { "shape": [1, 1, 16], "dtype": "float32", "dist": "uniform", "seed": 8105, "min": 0.1, "max": 0.9 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 1, 4096], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 16, 128, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-qwen3next-prefill-s512",
+      "preset": "model",
+      "provenance": { "notes": "Qwen3-Next class defaults over a 512-token prefill chunk." },
+      "vars": { "batch": 1, "seq": 512, "qHeads": 32, "kvHeads": 16, "dk": 128, "dv": 128 },
+      "attrs": {
+        "q_num_heads": 32,
+        "kv_num_heads": 16,
+        "update_rule": "gated_delta",
+        "scale": 0.08838834764831843,
+        "chunk_size": 64
+      },
+      "inputs": {
+        "queryT": { "shape": [1, 512, 4096], "dtype": "float32", "dist": "normal", "seed": 8200, "scale": 0.3 },
+        "keyT": { "shape": [1, 512, 2048], "dtype": "float32", "dist": "normal", "seed": 8201, "scale": 0.3 },
+        "valueT": { "shape": [1, 512, 2048], "dtype": "float32", "dist": "normal", "seed": 8202, "scale": 0.3 },
+        "pastStateT": { "shape": [1, 16, 128, 128], "dtype": "float32", "dist": "normal", "seed": 8203, "scale": 0.1 },
+        "decayT": { "shape": [1, 512, 2048], "dtype": "float32", "dist": "uniform", "seed": 8204, "min": 0.9, "max": 1 },
+        "betaT": { "shape": [1, 512, 16], "dtype": "float32", "dist": "uniform", "seed": 8205, "min": 0.1, "max": 0.9 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 512, 4096], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 16, 128, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    },
+    {
+      "name": "linear-attention-qwen3next-prefill-s2048",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-Next class defaults over a 2048-token prefill chunk, eight chunk_size 64 blocks per workgroup pass."
+      },
+      "vars": { "batch": 1, "seq": 2048, "qHeads": 32, "kvHeads": 16, "dk": 128, "dv": 128 },
+      "attrs": {
+        "q_num_heads": 32,
+        "kv_num_heads": 16,
+        "update_rule": "gated_delta",
+        "scale": 0.08838834764831843,
+        "chunk_size": 64
+      },
+      "inputs": {
+        "queryT": { "shape": [1, 2048, 4096], "dtype": "float32", "dist": "normal", "seed": 8300, "scale": 0.3 },
+        "keyT": { "shape": [1, 2048, 2048], "dtype": "float32", "dist": "normal", "seed": 8301, "scale": 0.3 },
+        "valueT": { "shape": [1, 2048, 2048], "dtype": "float32", "dist": "normal", "seed": 8302, "scale": 0.3 },
+        "pastStateT": { "shape": [1, 16, 128, 128], "dtype": "float32", "dist": "normal", "seed": 8303, "scale": 0.1 },
+        "decayT": {
+          "shape": [1, 2048, 2048],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 8304,
+          "min": 0.9,
+          "max": 1
+        },
+        "betaT": { "shape": [1, 2048, 16], "dtype": "float32", "dist": "uniform", "seed": 8305, "min": 0.1, "max": 0.9 }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 2048, 4096], "dtype": "float32" },
+        "presentStateT": { "shape": [1, 16, 128, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "gflops",
+            "value": "2 * args.batch * args.seq * (args.qHeads + 2 * args.kvHeads) * args.dk * args.dv"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/chunk-out.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,200 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}
+{% macro write_scalar(expr, dtype) %}
+{% if dtype == "float16" %}
+f16({{ expr }}){% else %}{{ expr }}{% endif %}
+{% endmacro -%}
+{% macro decay_at(bt, h, i) %}{% if decayPerElement %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }}) * HEAD_DIM_K + ({{ i }})]{% else %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }})]{% endif %}{%- endmacro %}
+{% macro emit_chunk_operands(needKtil=false, needQuery=false) %}
+{% if needQuery %}
+// The output scale defaults to 1/sqrt(head_dim_k), matching the recurrent kernels.
+fn out_scale() -> f32 {
+  return select(inverseSqrt(f32(HEAD_DIM_K)), params.scale, params.scale != 0.0);
+}
+{% endif %}
+fn k_hat(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw / {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% if needKtil %}
+fn k_til(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% endif %}
+{% if needQuery %}
+// The output scale rides on q_til so both output terms -- q_til * S and P * delta,
+// where P is itself built from q_til -- pick it up without a second pass over y.
+fn q_til(bt: u32, q_head: u32, {% if usesDecay %}h: u32, {% endif %}i: u32, scale: f32) -> f32 {
+  let raw = {{ read_scalar("query", "bt * params.qPackedDim + q_head * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }} * scale;
+{% else %}
+  return raw * scale;
+{% endif %}
+}
+{% endif %}
+{%- endmacro %}{% macro q_til_call(bt, q_head, h, i, scale) %}q_til({{ bt }}, {{ q_head }}, {% if usesDecay %}{{ h }}, {% endif %}{{ i }}, {{ scale }}){%- endmacro %}{% macro emit_chunk_head_setup(needHeads=false) %}
+{% if needHeads %}
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * HEAD_DIM_V;
+{% endif %}
+  let num_chunks = (params.seqLength + CHUNK - 1u) / CHUNK;
+{%- endmacro %}
+{% if queryDtype == "float16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.LinearAttention, chunked prefill: outputs, from entry states.
+//   y_t = q_til_t . S_c  +  sum_{s <= t} P[t, s] * delta_s,   P[t, s] = dot(k_hat_s, q_til_t)
+// S_c is the state entering this token's chunk, so once the sequential pass has
+// published one state per chunk, every chunk's output block is independent -- this
+// pass carries the largest share of the operator's arithmetic and runs entirely in
+// parallel. P is built here rather than stored: one workgroup owns a whole (output
+// head, chunk) block, so nothing recomputes it.
+//
+// A thread owns one value column and TOKEN_ROWS of the chunk, and the token block is
+// the outer loop, so nothing staged here grows with the chunk length. That is what
+// lets the chunk be long: the entry-state buffer the sequential pass publishes shrinks
+// as 1 / CHUNK, while this pass's re-read of it depends only on TOKEN_ROWS.
+const CHUNK: u32 = {{ chunkSize }}u;
+const HEAD_DIM_K: u32 = {{ headDimK }}u;
+const HEAD_DIM_V: u32 = {{ headDimV }}u;
+const KV_PER_KEY_HEAD: u32 = {{ kvPerKeyHead }}u;
+const TOKEN_ROWS: u32 = {{ chunkOutRows }}u;
+const TK: u32 = {{ chunkTileK }}u;
+const WG: u32 = {{ workgroupSize }}u;
+{{ emit_chunk_operands(needQuery=true) }}
+var<workgroup> qblock: array<f32, TOKEN_ROWS * HEAD_DIM_K>;
+var<workgroup> ktile: array<f32, CHUNK * TK>;
+var<workgroup> pmrow: array<f32, TOKEN_ROWS * CHUNK>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  let tid = lid.x;
+  let flat = wg.x + wg.y * nwg.x;
+{{ emit_chunk_head_setup(needHeads=true) }}
+  let out_heads = max(params.qNumHeads, params.kvNumHeads);
+  let chunk = flat % num_chunks;
+  let out_head = (flat / num_chunks) % out_heads;
+  let batch = flat / (num_chunks * out_heads);
+  if (batch >= params.batchSize) {
+    return;
+  }
+  let head = out_head / heads_per_group;
+  let q_head = (head * params.qNumHeads) / params.kvNumHeads + out_head % heads_per_group;
+  let base_t = batch * params.seqLength + chunk * CHUNK;
+  let live = min(CHUNK, params.seqLength - chunk * CHUNK);
+  let scale = out_scale();
+  let state_base = ((batch * params.kvNumHeads + head) * num_chunks + chunk) * HEAD_DIM_K * HEAD_DIM_V;
+{% if usesBeta %}
+  let chunk_base = ((batch * params.kvNumHeads + head) * num_chunks + chunk) * CHUNK;
+{% endif %}
+  for (var tb = 0u; tb < CHUNK; tb = tb + TOKEN_ROWS) {
+    for (var e = tid; e < TOKEN_ROWS * HEAD_DIM_K; e = e + WG) {
+      let t = tb + e / HEAD_DIM_K;
+      var qv = 0.0;
+      if (t < live) {
+        qv = {{ q_til_call("base_t + t", "q_head", "head", "e % HEAD_DIM_K", "scale") }};
+      }
+      qblock[e] = qv;
+    }
+    for (var e = tid; e < TOKEN_ROWS * CHUNK; e = e + WG) {
+      pmrow[e] = 0.0;
+    }
+    var acc: array<f32, TOKEN_ROWS>;
+    for (var m = 0u; m < TOKEN_ROWS; m = m + 1u) {
+      acc[m] = 0.0;
+    }
+    workgroupBarrier();
+    // One sweep of the reduction axis feeds both output terms: this block's rows of P,
+    // and its q_til * S contribution. The state column a thread loads is used by all
+    // TOKEN_ROWS of its accumulators.
+    for (var kb = 0u; kb < HEAD_DIM_K; kb = kb + TK) {
+      for (var e = tid; e < CHUNK * TK; e = e + WG) {
+        let s = e / TK;
+        let i = kb + e % TK;
+        var kv = 0.0;
+        if (s < live && i < HEAD_DIM_K) {
+          kv = k_hat(base_t + s, head, i);
+        }
+        ktile[e] = kv;
+      }
+      workgroupBarrier();
+      // P is inclusive-lower: token t reads the state after its own update, matching
+      // the recurrent kernels.
+      for (var e = tid; e < TOKEN_ROWS * CHUNK; e = e + WG) {
+        let m = e / CHUNK;
+        let s = e % CHUNK;
+        if (s <= tb + m) {
+          var total = 0.0;
+          for (var u = 0u; u < TK; u = u + 1u) {
+            if (kb + u < HEAD_DIM_K) {
+              total = total + qblock[m * HEAD_DIM_K + kb + u] * ktile[s * TK + u];
+            }
+          }
+          pmrow[e] = pmrow[e] + total;
+        }
+      }
+      if (tid < HEAD_DIM_V) {
+        for (var u = 0u; u < TK; u = u + 1u) {
+          let i = kb + u;
+          if (i < HEAD_DIM_K) {
+            let sv = states[state_base + i * HEAD_DIM_V + tid];
+            for (var m = 0u; m < TOKEN_ROWS; m = m + 1u) {
+              acc[m] = acc[m] + qblock[m * HEAD_DIM_K + i] * sv;
+            }
+          }
+        }
+      }
+      workgroupBarrier();
+    }
+    if (tid < HEAD_DIM_V) {
+      for (var s = 0u; s < CHUNK; s = s + 1u) {
+{% if usesBeta %}
+        let dv_val = deltas[(chunk_base + s) * HEAD_DIM_V + tid];
+{% else %}
+        var dv_val = 0.0;
+        if (s < live) {
+          dv_val = {{ read_scalar("value", "(base_t + s) * params.vPackedDim + head * HEAD_DIM_V + tid", queryDtype) }};
+        }
+{% endif %}
+        for (var m = 0u; m < TOKEN_ROWS; m = m + 1u) {
+          acc[m] = acc[m] + pmrow[m * CHUNK + s] * dv_val;
+        }
+      }
+      for (var m = 0u; m < TOKEN_ROWS; m = m + 1u) {
+        let t = tb + m;
+        if (t < live) {
+          output[(base_t + t) * packed_out + out_head * HEAD_DIM_V + tid] =
+            {{ write_scalar("acc[m]", queryDtype) }};
+        }
+      }
+    }
+    workgroupBarrier();
+  }
+}

build/webgpu/chunk-prep.wgsl.jinja ADDED Viewed

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

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}{% if queryDtype == "float16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.LinearAttention, chunked prefill: within-chunk decay prefix.
+//   gexp[b, t, c] = exp(sum of decay[b, s, c] for s from the chunk start through t)
+// Every later pass reads the recurrence's decay through this one buffer. Writing
+// exp(prefix) rather than the prefix itself keeps the per-element exponential out
+// of the O(chunk^2) inner loops, where it would cost one transcendental per
+// multiply-add instead of one per element.
+const CHUNK: u32 = {{ chunkSize }}u;
+const WG: u32 = {{ workgroupSize }}u;
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  // 2D-folded flat (batch * chunk) index: wg.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit.
+  let flat = wg.x + wg.y * nwg.x;
+  let num_chunks = (params.seqLength + CHUNK - 1u) / CHUNK;
+  let chunk = flat % num_chunks;
+  let batch = flat / num_chunks;
+  if (batch >= params.batchSize) {
+    return;
+  }
+  let t0 = chunk * CHUNK;
+  let t1 = min(t0 + CHUNK, params.seqLength);
+  let packed = params.decayPackedDim;
+  // One thread owns a decay column and walks the chunk in order: the prefix is
+  // serial in t but independent across columns, so the whole chunk grid runs at once.
+  for (var col = lid.x; col < packed; col = col + WG) {
+    var prefix = 0.0;
+    for (var t = t0; t < t1; t = t + 1u) {
+      let idx = (batch * params.seqLength + t) * packed + col;
+      prefix = prefix + {{ read_scalar("decay", "idx", queryDtype) }};
+      gexp[idx] = exp(prefix);
+    }
+  }
+}

build/webgpu/chunk-scan.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,216 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}
+{% macro write_scalar(expr, dtype) %}
+{% if dtype == "float16" %}
+f16({{ expr }}){% else %}{{ expr }}{% endif %}
+{% endmacro -%}
+{% macro decay_at(bt, h, i) %}{% if decayPerElement %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }}) * HEAD_DIM_K + ({{ i }})]{% else %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }})]{% endif %}{%- endmacro %}
+{% macro emit_chunk_operands(needKtil=false, needQuery=false) %}
+{% if needQuery %}
+// The output scale defaults to 1/sqrt(head_dim_k), matching the recurrent kernels.
+fn out_scale() -> f32 {
+  return select(inverseSqrt(f32(HEAD_DIM_K)), params.scale, params.scale != 0.0);
+}
+{% endif %}
+fn k_hat(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw / {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% if needKtil %}
+fn k_til(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% endif %}
+{% if needQuery %}
+// The output scale rides on q_til so both output terms -- q_til * S and P * delta,
+// where P is itself built from q_til -- pick it up without a second pass over y.
+fn q_til(bt: u32, q_head: u32, {% if usesDecay %}h: u32, {% endif %}i: u32, scale: f32) -> f32 {
+  let raw = {{ read_scalar("query", "bt * params.qPackedDim + q_head * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }} * scale;
+{% else %}
+  return raw * scale;
+{% endif %}
+}
+{% endif %}
+{%- endmacro %}{% macro emit_chunk_head_setup(needHeads=false) %}
+{% if needHeads %}
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * HEAD_DIM_V;
+{% endif %}
+  let num_chunks = (params.seqLength + CHUNK - 1u) / CHUNK;
+{%- endmacro %}
+{% if queryDtype == "float16" or stateDtype == "float16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.LinearAttention, chunked prefill: the sequential state scan.
+//
+// This is the only pass that must run chunk by chunk, and it is deliberately the
+// smallest: two dense products per chunk rather than one dependent step per token.
+//   delta_c = u_c - wk_c S_c            (identity for the non-delta rules, where delta = v)
+//   S_{c+1} = A_c * (S_c + K_hat_c^T delta_c)
+// It publishes each chunk's entry state and correction so the output pass, which
+// carries most of the arithmetic, can run over every chunk at once.
+const CHUNK: u32 = {{ chunkSize }}u;
+const HEAD_DIM_K: u32 = {{ headDimK }}u;
+const HEAD_DIM_V: u32 = {{ headDimV }}u;
+const KV_PER_KEY_HEAD: u32 = {{ kvPerKeyHead }}u;
+const TILE_V: u32 = {{ chunkTileV }}u;
+const TOKEN_TILE: u32 = {{ chunkScanTokens }}u;
+const WG: u32 = {{ workgroupSize }}u;
+// One thread per (value column, group). The two phases split the workgroup along
+// different axes -- tokens for the correction, reduction rows for the state update --
+// so neither needs a cross-thread fold.
+const GROUPS: u32 = WG / TILE_V;
+{% if usesBeta %}
+const TOKENS_PER_TILE: u32 = TOKEN_TILE / GROUPS;
+{% endif %}
+const ROWS_PER_GROUP: u32 = HEAD_DIM_K / GROUPS;
+{{ emit_chunk_operands() }}
+var<workgroup> st: array<f32, HEAD_DIM_K * TILE_V>;
+var<workgroup> dl: array<f32, CHUNK * TILE_V>;
+// Staging for TOKEN_TILE rows of wk, then of k_hat: same shape, disjoint live ranges.
+var<workgroup> stage: array<f32, TOKEN_TILE * HEAD_DIM_K>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  let tid = lid.x;
+  let col = tid % TILE_V;
+  let group = tid / TILE_V;
+  let flat = wg.x + wg.y * nwg.x;
+{{ emit_chunk_head_setup() }}
+  let v_tiles = HEAD_DIM_V / TILE_V;
+  let v_tile = flat % v_tiles;
+  let head = (flat / v_tiles) % params.kvNumHeads;
+  let batch = flat / (v_tiles * params.kvNumHeads);
+  if (batch >= params.batchSize) {
+    return;
+  }
+  let dv0 = v_tile * TILE_V;
+  let head_base = (batch * params.kvNumHeads + head) * HEAD_DIM_K;
+  for (var e = tid; e < HEAD_DIM_K * TILE_V; e = e + WG) {
+{% if hasPastState %}
+    st[e] = {{ read_scalar("past_state", "(head_base + e / TILE_V) * HEAD_DIM_V + dv0 + e % TILE_V", stateDtype) }};
+{% else %}
+    st[e] = 0.0;
+{% endif %}
+  }
+  workgroupBarrier();
+  for (var chunk = 0u; chunk < num_chunks; chunk = chunk + 1u) {
+    let base_t = batch * params.seqLength + chunk * CHUNK;
+    let live = min(CHUNK, params.seqLength - chunk * CHUNK);
+{% if usesBeta %}
+    let chunk_base = ((batch * params.kvNumHeads + head) * num_chunks + chunk) * CHUNK;
+{% endif %}
+{% if usesBeta %}
+    // delta = u - wk S. Staging wk by token tile keeps the reduction loop entirely in
+    // workgroup memory without holding the whole chunk.
+    for (var tb = 0u; tb < CHUNK; tb = tb + TOKEN_TILE) {
+      for (var e = tid; e < TOKEN_TILE * HEAD_DIM_K; e = e + WG) {
+        stage[e] = wk[(chunk_base + tb + e / HEAD_DIM_K) * HEAD_DIM_K + e % HEAD_DIM_K];
+      }
+      workgroupBarrier();
+      for (var m = 0u; m < TOKENS_PER_TILE; m = m + 1u) {
+        let local_t = group * TOKENS_PER_TILE + m;
+        let t = tb + local_t;
+        var a = uvec[(chunk_base + t) * HEAD_DIM_V + dv0 + col];
+        for (var i = 0u; i < HEAD_DIM_K; i = i + 1u) {
+          a = a - stage[local_t * HEAD_DIM_K + i] * st[i * TILE_V + col];
+        }
+        dl[t * TILE_V + col] = a;
+      }
+      workgroupBarrier();
+    }
+    for (var e = tid; e < CHUNK * TILE_V; e = e + WG) {
+      deltas[(chunk_base + e / TILE_V) * HEAD_DIM_V + dv0 + e % TILE_V] = dl[e];
+    }
+{% else %}
+    // The non-delta rules take delta = v with no correction at all.
+    for (var e = tid; e < CHUNK * TILE_V; e = e + WG) {
+      let t = e / TILE_V;
+      var v = 0.0;
+      if (t < live) {
+        v = {{ read_scalar("value", "(base_t + t) * params.vPackedDim + head * HEAD_DIM_V + dv0 + e % TILE_V", queryDtype) }};
+      }
+      dl[e] = v;
+    }
+{% endif %}
+    // Publish the entry state before the update: it is what the output pass reads.
+    let state_base = ((batch * params.kvNumHeads + head) * num_chunks + chunk) * HEAD_DIM_K * HEAD_DIM_V;
+    for (var e = tid; e < HEAD_DIM_K * TILE_V; e = e + WG) {
+      states[state_base + (e / TILE_V) * HEAD_DIM_V + dv0 + e % TILE_V] = st[e];
+    }
+    workgroupBarrier();
+    // S = A_c * (S + K_hat^T delta): an outer-product accumulation, so splitting the
+    // workgroup along the reduction axis here keeps every thread's rows private. The
+    // row accumulators live in registers across the whole chunk, which is what makes
+    // each staged k_hat element feed ROWS_PER_GROUP multiply-adds instead of one.
+    var acc: array<f32, ROWS_PER_GROUP>;
+    for (var m = 0u; m < ROWS_PER_GROUP; m = m + 1u) {
+      acc[m] = st[(group * ROWS_PER_GROUP + m) * TILE_V + col];
+    }
+    for (var tb = 0u; tb < CHUNK; tb = tb + TOKEN_TILE) {
+      for (var e = tid; e < TOKEN_TILE * HEAD_DIM_K; e = e + WG) {
+        let t = tb + e / HEAD_DIM_K;
+        var v = 0.0;
+        if (t < live) {
+          v = k_hat(base_t + t, head, e % HEAD_DIM_K);
+        }
+        stage[e] = v;
+      }
+      workgroupBarrier();
+      for (var lt = 0u; lt < TOKEN_TILE; lt = lt + 1u) {
+        let dv_val = dl[(tb + lt) * TILE_V + col];
+        for (var m = 0u; m < ROWS_PER_GROUP; m = m + 1u) {
+          acc[m] = acc[m] + stage[lt * HEAD_DIM_K + group * ROWS_PER_GROUP + m] * dv_val;
+        }
+      }
+      workgroupBarrier();
+    }
+{% if usesDecay %}
+    let last_t = base_t + live - 1u;
+{% endif %}
+    for (var m = 0u; m < ROWS_PER_GROUP; m = m + 1u) {
+      let i = group * ROWS_PER_GROUP + m;
+{% if usesDecay %}
+      st[i * TILE_V + col] = acc[m] * {{ decay_at("last_t", "head", "i") }};
+{% else %}
+      st[i * TILE_V + col] = acc[m];
+{% endif %}
+    }
+    workgroupBarrier();
+  }
+  for (var e = tid; e < HEAD_DIM_K * TILE_V; e = e + WG) {
+    present_state[(head_base + e / TILE_V) * HEAD_DIM_V + dv0 + e % TILE_V] =
+      {{ write_scalar("st[e]", stateDtype) }};
+  }
+}

build/webgpu/chunk-ut.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,234 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}{% macro decay_at(bt, h, i) %}{% if decayPerElement %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }}) * HEAD_DIM_K + ({{ i }})]{% else %}gexp[({{ bt }}) * params.decayPackedDim + ({{ h }})]{% endif %}{%- endmacro %}
+{% macro emit_chunk_operands(needKtil=false, needQuery=false) %}
+{% if needQuery %}
+// The output scale defaults to 1/sqrt(head_dim_k), matching the recurrent kernels.
+fn out_scale() -> f32 {
+  return select(inverseSqrt(f32(HEAD_DIM_K)), params.scale, params.scale != 0.0);
+}
+{% endif %}
+fn k_hat(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw / {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% if needKtil %}
+fn k_til(bt: u32, h: u32, i: u32) -> f32 {
+  let raw = {{ read_scalar("key", "bt * params.kPackedDim + (h / KV_PER_KEY_HEAD) * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }};
+{% else %}
+  return raw;
+{% endif %}
+}
+{% endif %}
+{% if needQuery %}
+// The output scale rides on q_til so both output terms -- q_til * S and P * delta,
+// where P is itself built from q_til -- pick it up without a second pass over y.
+fn q_til(bt: u32, q_head: u32, {% if usesDecay %}h: u32, {% endif %}i: u32, scale: f32) -> f32 {
+  let raw = {{ read_scalar("query", "bt * params.qPackedDim + q_head * HEAD_DIM_K + i", queryDtype) }};
+{% if usesDecay %}
+  return raw * {{ decay_at("bt", "h", "i") }} * scale;
+{% else %}
+  return raw * scale;
+{% endif %}
+}
+{% endif %}
+{%- endmacro %}{% macro emit_chunk_head_setup(needHeads=false) %}
+{% if needHeads %}
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * HEAD_DIM_V;
+{% endif %}
+  let num_chunks = (params.seqLength + CHUNK - 1u) / CHUNK;
+{%- endmacro %}
+{% if queryDtype == "float16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.LinearAttention, chunked prefill: the chunk-local delta transform.
+//
+// The delta rules apply S_t = (I - beta_t k_t k_t^T) diag(a_t) S_{t-1} + beta_t k_t v_t^T,
+// whose per-token correction is what forces the recurrent kernels to serialize. Over
+// one chunk the corrections satisfy a unit lower triangular system
+//   (I + W) d = beta * (v - S_0^T k_til),   W[t,s] = beta_t * dot(k_til_t, k_hat_s), s < t
+// so d = u - wk * S_0 with u = Tb * V and wk = Tb * K_til, Tb = (I + W)^-1 diag(beta).
+// Both depend only on this chunk's own key, value and beta -- not on the entry state --
+// so every chunk in the sequence computes them at once here, and the sequential pass is
+// left with two dense products per chunk instead of one dependent step per token.
+const CHUNK: u32 = {{ chunkSize }}u;
+const HEAD_DIM_K: u32 = {{ headDimK }}u;
+const HEAD_DIM_V: u32 = {{ headDimV }}u;
+const KV_PER_KEY_HEAD: u32 = {{ kvPerKeyHead }}u;
+const TK: u32 = {{ chunkTileK }}u;
+const WG: u32 = {{ workgroupSize }}u;
+// Entries of the CHUNK x CHUNK system each thread carries across the reduction tiles.
+const ENTRIES: u32 = (CHUNK * CHUNK) / WG;
+{{ emit_chunk_operands(needKtil=true) }}
+// `wm` holds W while the system is being built, then the inverse in place: row t of
+// the inverse is produced from row t of W and the already-final rows above it.
+var<workgroup> wm: array<f32, CHUNK * CHUNK>;
+var<workgroup> rowbuf: array<f32, CHUNK>;
+var<workgroup> betas: array<f32, CHUNK>;
+var<workgroup> tile: array<f32, CHUNK * TK>;
+// Both W operands are staged: the inner product below runs CHUNK^2 times per tile,
+// so a global read there would be paid once per pair rather than once per element.
+var<workgroup> tile_til: array<f32, CHUNK * TK>;
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  let tid = lid.x;
+  let flat = wg.x + wg.y * nwg.x;
+{{ emit_chunk_head_setup() }}
+  let chunk = flat % num_chunks;
+  let head = (flat / num_chunks) % params.kvNumHeads;
+  let batch = flat / (num_chunks * params.kvNumHeads);
+  if (batch >= params.batchSize) {
+    return;
+  }
+  let base_t = batch * params.seqLength + chunk * CHUNK;
+  // Tokens past the sequence end carry beta 0, which makes their rows of the system
+  // the identity and their outputs zero, so the tail chunk needs no separate arm.
+  let live = min(CHUNK, params.seqLength - chunk * CHUNK);
+  if (tid < CHUNK) {
+    var b = 0.0;
+    if (tid < live) {
+      let bt = base_t + tid;
+      b = {{ read_scalar("beta", "select(bt * params.kvNumHeads + head, bt, params.betaPackedDim == 1u)", queryDtype) }};
+    }
+    betas[tid] = b;
+  }
+  workgroupBarrier();
+  // W = beta * tril(K_til K_hat^T, -1), accumulated over reduction-axis tiles so the
+  // two operands share one staging buffer pass instead of a full CHUNK x HEAD_DIM_K copy.
+  var acc: array<f32, ENTRIES>;
+  for (var m = 0u; m < ENTRIES; m = m + 1u) {
+    acc[m] = 0.0;
+  }
+  for (var kb = 0u; kb < HEAD_DIM_K; kb = kb + TK) {
+    for (var e = tid; e < CHUNK * TK; e = e + WG) {
+      let t = e / TK;
+      let i = kb + e % TK;
+      var hat = 0.0;
+      var til = 0.0;
+      if (t < live && i < HEAD_DIM_K) {
+        hat = k_hat(base_t + t, head, i);
+        til = k_til(base_t + t, head, i);
+      }
+      tile[e] = hat;
+      tile_til[e] = til;
+    }
+    workgroupBarrier();
+    for (var m = 0u; m < ENTRIES; m = m + 1u) {
+      let e = tid + m * WG;
+      let t = e / CHUNK;
+      let s = e % CHUNK;
+      if (s < t) {
+        var total = 0.0;
+        for (var u = 0u; u < TK; u = u + 1u) {
+          total = total + tile_til[t * TK + u] * tile[s * TK + u];
+        }
+        acc[m] = acc[m] + total;
+      }
+    }
+    workgroupBarrier();
+  }
+  for (var m = 0u; m < ENTRIES; m = m + 1u) {
+    let e = tid + m * WG;
+    wm[e] = acc[m] * betas[e / CHUNK];
+  }
+  workgroupBarrier();
+  // Forward substitution for (I + W)^-1, one row per step. W is strictly lower, so
+  // row t reads only finished rows; `rowbuf` copies row t out before it is overwritten.
+  for (var t = 0u; t < CHUNK; t = t + 1u) {
+    if (tid < CHUNK) {
+      rowbuf[tid] = wm[t * CHUNK + tid];
+    }
+    workgroupBarrier();
+    if (tid < CHUNK) {
+      var a = select(0.0, 1.0, tid == t);
+      for (var r = 0u; r < t; r = r + 1u) {
+        a = a - rowbuf[r] * wm[r * CHUNK + tid];
+      }
+      wm[t * CHUNK + tid] = a;
+    }
+    workgroupBarrier();
+  }
+  // Tb = T diag(beta): scaling by column completes the transform.
+  for (var e = tid; e < CHUNK * CHUNK; e = e + WG) {
+    wm[e] = wm[e] * betas[e % CHUNK];
+  }
+  workgroupBarrier();
+  let chunk_base = ((batch * params.kvNumHeads + head) * num_chunks + chunk) * CHUNK;
+  // wk = Tb K_til, staged one reduction tile at a time.
+  for (var kb = 0u; kb < HEAD_DIM_K; kb = kb + TK) {
+    for (var e = tid; e < CHUNK * TK; e = e + WG) {
+      let s = e / TK;
+      let i = kb + e % TK;
+      var v = 0.0;
+      if (s < live && i < HEAD_DIM_K) {
+        v = k_til(base_t + s, head, i);
+      }
+      tile[e] = v;
+    }
+    workgroupBarrier();
+    for (var e = tid; e < CHUNK * TK; e = e + WG) {
+      let t = e / TK;
+      let u = e % TK;
+      if (kb + u < HEAD_DIM_K) {
+        var total = 0.0;
+        for (var s = 0u; s <= t; s = s + 1u) {
+          total = total + wm[t * CHUNK + s] * tile[s * TK + u];
+        }
+        wk[(chunk_base + t) * HEAD_DIM_K + kb + u] = total;
+      }
+    }
+    workgroupBarrier();
+  }
+  // u = Tb V, over the value axis.
+  for (var vb = 0u; vb < HEAD_DIM_V; vb = vb + TK) {
+    for (var e = tid; e < CHUNK * TK; e = e + WG) {
+      let s = e / TK;
+      let j = vb + e % TK;
+      var v = 0.0;
+      if (s < live && j < HEAD_DIM_V) {
+        v = {{ read_scalar("value", "(base_t + s) * params.vPackedDim + head * HEAD_DIM_V + j", queryDtype) }};
+      }
+      tile[e] = v;
+    }
+    workgroupBarrier();
+    for (var e = tid; e < CHUNK * TK; e = e + WG) {
+      let t = e / TK;
+      let u = e % TK;
+      if (vb + u < HEAD_DIM_V) {
+        var total = 0.0;
+        for (var s = 0u; s <= t; s = s + 1u) {
+          total = total + wm[t * CHUNK + s] * tile[s * TK + u];
+        }
+        uvec[(chunk_base + t) * HEAD_DIM_V + vb + u] = total;
+      }
+    }
+    workgroupBarrier();
+  }
+}

build/webgpu/linear-attention.scalar.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,332 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}
+{% macro write_scalar(expr, dtype) %}
+{% if dtype == "float16" %}
+f16({{ expr }}){% else %}{{ expr }}{% endif %}
+{% endmacro -%}
+{% macro emit_tiled_setup(dvGroups=1) %}
+  let head_dim_k = params.qPackedDim / params.qNumHeads;
+  let head_dim_v = params.vPackedDim / params.kvNumHeads;
+  let n_key_heads = params.kPackedDim / head_dim_k;
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let kv_per_key_head = params.kvNumHeads / n_key_heads;
+{% if dvGroups == 1 %}
+  let dv_tiles = (head_dim_v + TILE_V - 1u) / TILE_V;
+{% else %}
+  // Tile slots per workgroup-index step: each step covers DV_GROUPS value tiles.
+  let dv_tiles = ((head_dim_v + TILE_V - 1u) / TILE_V + DV_GROUPS - 1u) / DV_GROUPS;
+{% endif %}
+  let scale = select(inverseSqrt(f32(head_dim_k)), params.scale, params.scale != 0.0);
+  // 2D-folded flat (batch*head*dv_tile) index: wg.y carries the high bits past
+  // the maxComputeWorkgroupsPerDimension dispatch limit. Reduces to wg.x when nwg.y == 1; the batch_idx >=
+  // params.batchSize guard drops the over-dispatched tail.
+  let workgroup_idx = wg.x + wg.y * nwg.x;
+  let dv_tile_idx = workgroup_idx % dv_tiles;
+  let bh = workgroup_idx / dv_tiles;
+  let head_idx = bh % params.kvNumHeads;
+  let batch_idx = bh / params.kvNumHeads;
+  if (batch_idx >= params.batchSize) {
+    return;
+  }
+{% if dvGroups == 1 %}
+  let dv_start = dv_tile_idx * TILE_V;
+{% else %}
+  let dv_start = (dv_tile_idx * DV_GROUPS + dv_group) * TILE_V;
+{% endif %}
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * head_dim_v;
+  let key_head_idx = head_idx / kv_per_key_head;
+{%- endmacro -%}
+{% macro emit_query_groups(first_group) %}
+    for (var qg = {{ first_group }}u; qg < heads_per_group; qg = qg + 1u) {
+      let q_head = (head_idx * params.qNumHeads) / params.kvNumHeads + qg;
+      let out_head = head_idx * heads_per_group + qg;
+      var q_val = 0.0;
+{% if not useSubgroups %}
+      var local_pre: array<f32, TILE_V>;
+{% endif %}
+      if (tid < head_dim_k) {
+        let q_idx = bt * params.qPackedDim + q_head * head_dim_k + tid;
+        q_val = {{ read_scalar("query", "q_idx", queryDtype) }};
+      }
+{% if useSubgroups %}
+      let qg_subgroup_index = tid / sg_size;
+      let qg_subgroup_count = (WG + sg_size - 1u) / sg_size;
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        let sg_pre = subgroupAdd(state[j] * q_val);
+        if (sg_lid == 0u) {
+          red_preout[j * WG + qg_subgroup_index] = sg_pre;
+        }
+      }
+      workgroupBarrier();
+      if (WG > sg_size) {
+        if (tid < TILE_V) {
+          var pre_total = 0.0;
+          for (var i = 1u; i < qg_subgroup_count; i = i + 1u) {
+            pre_total = pre_total + red_preout[tid * WG + i];
+          }
+          red_preout[tid * WG] = red_preout[tid * WG] + pre_total;
+        }
+        workgroupBarrier();
+      }
+{% else %}
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        red_preout[j * WG + tid] = state[j] * q_val;
+      }
+      workgroupBarrier();
+      for (var j = tid; j < TILE_V; j = j + WG) {
+        var pre_total = red_preout[j * WG];
+        for (var i = 1u; i < WG; i = i + 1u) {
+          pre_total = pre_total + red_preout[j * WG + i];
+        }
+        local_pre[j] = pre_total;
+      }
+      workgroupBarrier();
+{% endif %}
+{% if useSubgroups %}
+      if (tid == 0u) {
+        let out_base = bt * packed_out + out_head * head_dim_v + dv_start;
+        for (var j = 0u; j < TILE_V; j = j + 1u) {
+          if (dv_start + j < head_dim_v) {
+            output[out_base + j] = {{ write_scalar("red_preout[j * WG] * scale", outputDtype) }};
+          }
+        }
+      }
+      workgroupBarrier();
+{% else %}
+      let out_base = bt * packed_out + out_head * head_dim_v + dv_start;
+      for (var j = tid; j < TILE_V; j = j + WG) {
+        if (dv_start + j < head_dim_v) {
+          output[out_base + j] = {{ write_scalar("local_pre[j] * scale", outputDtype) }};
+        }
+      }
+{% endif %}
+    }
+{%- endmacro %}
+{% set usesDecay = updateRule == "gated" or updateRule == "gated_delta" %}
+{% set usesBeta = updateRule == "delta" or updateRule == "gated_delta" %}
+{% if queryDtype == "float16" or stateDtype == "float16" %}
+enable f16;
+{% endif %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+const TILE_V: u32 = {{ tileV }}u;
+{% if usesBeta %}var<workgroup> red_retrieved: array<f32, WG * TILE_V>;
+{% endif %}var<workgroup> red_preout: array<f32, WG * TILE_V>;
+{% if usesBeta %}var<workgroup> red_kq: array<f32, WG>;
+var<workgroup> broadcast_delta: array<f32, TILE_V>;
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+  @builtin(subgroup_invocation_id) sg_lid: u32,
+  @builtin(subgroup_size) sg_size: u32{% endif %}
+) {
+  let tid = lid.x;
+{{ emit_tiled_setup() }}
+  var state: array<f32, TILE_V>;
+  for (var j = 0u; j < TILE_V; j = j + 1u) {
+    state[j] = 0.0;
+  }
+{% if hasPastState %}
+  if (tid < head_dim_k) {
+{% if hasStateWindow %}
+    // A windowed past_state is read only from slot stateWindow-1, the state after
+    // the last token of the previous call.
+{% endif %}
+    let state_base = {% if hasStateWindow %}(params.stateWindow - 1u) * params.stateSlotStride + {% endif %}((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + tid) * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        state[j] = {{ read_scalar("past_state", "state_base + j", stateDtype) }};
+      }
+    }
+  }
+{% endif %}
+{% if hasStateWindow %}
+  // Slots below max(0, stateWindow - seqLength) hold no token from this call.
+  if (tid < head_dim_k) {
+    for (var z = 0u; z + params.seqLength < params.stateWindow; z = z + 1u) {
+      let z_base = z * params.stateSlotStride
+        + ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + tid) * head_dim_v + dv_start;
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        if (dv_start + j < head_dim_v) {
+          present_state[z_base + j] = {{ write_scalar("0.0", stateDtype) }};
+        }
+      }
+    }
+  }
+{% endif %}
+  for (var t = 0u; t < params.seqLength; t = t + 1u) {
+    let bt = batch_idx * params.seqLength + t;
+    var k_val = 0.0;
+    if (tid < head_dim_k) {
+      let k_idx = bt * params.kPackedDim + key_head_idx * head_dim_k + tid;
+      k_val = {{ read_scalar("key", "k_idx", keyDtype) }};
+    }
+{% if usesDecay %}
+    var decay_factor = 1.0;
+    if (params.decayPackedDim == params.kvNumHeads) {
+      decay_factor = exp({{ read_scalar("decay", "bt * params.kvNumHeads + head_idx", decayDtype) }});
+    } else if (tid < head_dim_k) {
+      let decay_idx = bt * params.decayPackedDim + head_idx * head_dim_k + tid;
+      decay_factor = exp({{ read_scalar("decay", "decay_idx", decayDtype) }});
+    }
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      state[j] = state[j] * decay_factor;
+    }
+{% endif %}
+{% if usesBeta %}
+    let q_head_0 = (head_idx * params.qNumHeads) / params.kvNumHeads;
+    let out_head_0 = head_idx * heads_per_group;
+    var q0_val = 0.0;
+    if (tid < head_dim_k) {
+      let q0_idx = bt * params.qPackedDim + q_head_0 * head_dim_k + tid;
+      q0_val = {{ read_scalar("query", "q0_idx", queryDtype) }};
+    }
+{% if useSubgroups %}
+    let subgroup_index_b = tid / sg_size;
+    let subgroup_count_b = (WG + sg_size - 1u) / sg_size;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      let sg_rp = subgroupAdd(vec2<f32>(state[j] * k_val, state[j] * q0_val));
+      if (sg_lid == 0u) {
+        red_retrieved[j * WG + subgroup_index_b] = sg_rp.x;
+        red_preout[j * WG + subgroup_index_b] = sg_rp.y;
+      }
+    }
+    let sg_kq = subgroupAdd(k_val * q0_val);
+    if (sg_lid == 0u) {
+      red_kq[subgroup_index_b] = sg_kq;
+    }
+    workgroupBarrier();
+    if (WG > sg_size) {
+      if (tid < TILE_V) {
+        var ret_total = 0.0;
+        var pre_total = 0.0;
+        for (var i = 1u; i < subgroup_count_b; i = i + 1u) {
+          ret_total = ret_total + red_retrieved[tid * WG + i];
+          pre_total = pre_total + red_preout[tid * WG + i];
+        }
+        red_retrieved[tid * WG] = red_retrieved[tid * WG] + ret_total;
+        red_preout[tid * WG] = red_preout[tid * WG] + pre_total;
+      }
+      if (tid == TILE_V) {
+        var kq_total = 0.0;
+        for (var i = 1u; i < subgroup_count_b; i = i + 1u) {
+          kq_total = kq_total + red_kq[i];
+        }
+        red_kq[0] = red_kq[0] + kq_total;
+      }
+      workgroupBarrier();
+    }
+{% else %}
+    // Only the TILE_V output lanes consume these dot products. After one rendezvous,
+    // one lane folds each column and publishes it before the second barrier.
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      red_retrieved[j * WG + tid] = state[j] * k_val;
+      red_preout[j * WG + tid] = state[j] * q0_val;
+    }
+    red_kq[tid] = k_val * q0_val;
+    workgroupBarrier();
+    for (var j = tid; j < TILE_V; j = j + WG) {
+      var ret_total = red_retrieved[j * WG];
+      var pre_total = red_preout[j * WG];
+      for (var i = 1u; i < WG; i = i + 1u) {
+        ret_total = ret_total + red_retrieved[j * WG + i];
+        pre_total = pre_total + red_preout[j * WG + i];
+      }
+      red_retrieved[j * WG] = ret_total;
+      red_preout[j * WG] = pre_total;
+    }
+    if (tid == 0u) {
+      var kq_total = red_kq[0];
+      for (var i = 1u; i < WG; i = i + 1u) {
+        kq_total = kq_total + red_kq[i];
+      }
+      red_kq[0] = kq_total;
+    }
+    workgroupBarrier();
+{% endif %}
+    if (tid == 0u) {
+      var beta_idx = bt * params.kvNumHeads + head_idx;
+      if (params.betaPackedDim == 1u) {
+        beta_idx = bt;
+      }
+      let beta_val = {{ read_scalar("beta", "beta_idx", betaDtype) }};
+      let v_base = bt * params.vPackedDim + head_idx * head_dim_v + dv_start;
+      let out_base = bt * packed_out + out_head_0 * head_dim_v + dv_start;
+      let kq_dot = red_kq[0];
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        if (dv_start + j < head_dim_v) {
+          let v_val = {{ read_scalar("value", "v_base + j", valueDtype) }};
+          let delta_j = beta_val * (v_val - red_retrieved[j * WG]);
+          broadcast_delta[j] = delta_j;
+          output[out_base + j] = {{ write_scalar("(red_preout[j * WG] + delta_j * kq_dot) * scale", outputDtype) }};
+        } else {
+          broadcast_delta[j] = 0.0;
+        }
+      }
+    }
+    workgroupBarrier();
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      state[j] = state[j] + k_val * broadcast_delta[j];
+    }
+{{ emit_query_groups(1) }}
+{% else %}
+    let v_base = bt * params.vPackedDim + head_idx * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        state[j] = state[j] + k_val * {{ read_scalar("value", "v_base + j", valueDtype) }};
+      }
+    }
+{{ emit_query_groups(0) }}
+{% endif %}
+{% if hasStateWindow %}
+    // Snapshot the state after this token into its window slot. Slot j holds the
+    // state after token (seqLength - stateWindow + j), so this token owns slot
+    // (t + stateWindow - seqLength) whenever that lands inside the window.
+    if (t + params.stateWindow >= params.seqLength && tid < head_dim_k) {
+      let win_slot = t + params.stateWindow - params.seqLength;
+      let win_base = win_slot * params.stateSlotStride
+        + ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + tid) * head_dim_v + dv_start;
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        if (dv_start + j < head_dim_v) {
+          present_state[win_base + j] = {{ write_scalar("state[j]", stateDtype) }};
+        }
+      }
+    }
+{% endif %}
+  }
+{% if not hasStateWindow %}
+  if (tid < head_dim_k) {
+    let state_base = ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + tid) * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        present_state[state_base + j] = {{ write_scalar("state[j]", stateDtype) }};
+      }
+    }
+  }
+{% endif %}
+}

build/webgpu/linear-attention.serial.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,155 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}
+{% macro write_scalar(expr, dtype) %}
+{% if dtype == "float16" %}
+f16({{ expr }}){% else %}{{ expr }}{% endif %}
+{% endmacro -%}
+{% macro emit_serial_query_groups(first_group) %}
+    for (var qg = {{ first_group }}u; qg < heads_per_group; qg = qg + 1u) {
+      let q_head = (head_idx * params.qNumHeads) / params.kvNumHeads + qg;
+      let out_head = head_idx * heads_per_group + qg;
+      let q_base = bt * params.qPackedDim + q_head * HEAD_DIM_K;
+      var q_out = 0.0;
+      for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+        q_out = q_out + state[d] * {{ read_scalar("query", "q_base + d", queryDtype) }};
+      }
+      let out_idx = bt * packed_out + out_head * head_dim_v + dv_idx;
+      output[out_idx] = {{ write_scalar("q_out * scale", outputDtype) }};
+    }
+{%- endmacro -%}
+{% set gatedDeltaRule = updateRule == "gated_delta" %}
+{% if queryDtype == "float16" or stateDtype == "float16" %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const HEAD_DIM_K: u32 = {{ headDimK }}u;
+// Barrier-free small-dk recurrence. One invocation owns one
+// (batch, kv-head, value-dimension) state column and keeps its complete dk
+// slice private across the sequence. This trades dk-lane parallelism for zero
+// workgroup synchronization. Selection limits this route to small dk.
+@compute @workgroup_size(1, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+) {
+  let flat_idx = wg.x + wg.y * nwg.x;
+  let head_dim_v = params.vPackedDim / params.kvNumHeads;
+  let dv_idx = flat_idx % head_dim_v;
+  let bh = flat_idx / head_dim_v;
+  let head_idx = bh % params.kvNumHeads;
+  let batch_idx = bh / params.kvNumHeads;
+  if (batch_idx >= params.batchSize) {
+    return;
+  }
+  let n_key_heads = params.kPackedDim / HEAD_DIM_K;
+  let kv_per_key_head = params.kvNumHeads / n_key_heads;
+  let key_head_idx = head_idx / kv_per_key_head;
+  // Standard GQA has qNumHeads >= kvNumHeads and emits one output head per query head. Inverse
+  // GQA has kvNumHeads > qNumHeads and emits one per KV head, with several KV heads sharing a
+  // query head. max() makes the group count 1
+  // in that case, and then these two formulas cover both layouts with no branch:
+  //   q_head   = (head_idx * qNumHeads) / kvNumHeads + qg
+  //   out_head = head_idx * heads_per_group + qg
+  // In the standard layout the division is exact and the two agree; in the inverse layout the
+  // group count is 1, so q_head floors several KV heads onto one query head and out_head is the
+  // KV head itself.
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * head_dim_v;
+  let scale = select(inverseSqrt(f32(HEAD_DIM_K)), params.scale, params.scale != 0.0);
+  var state: array<f32, HEAD_DIM_K>;
+  for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+{% if hasPastState %}
+    let state_idx = {% if hasStateWindow %}(params.stateWindow - 1u) * params.stateSlotStride + {% endif %}((batch_idx * params.kvNumHeads + head_idx) * HEAD_DIM_K + d) * head_dim_v + dv_idx;
+    state[d] = {{ read_scalar("past_state", "state_idx", stateDtype) }};
+{% else %}
+    state[d] = 0.0;
+{% endif %}
+  }
+  var key_local: array<f32, HEAD_DIM_K>;
+{% if hasStateWindow %}
+  // Slots below max(0, stateWindow - seqLength) hold no token from this call.
+  for (var z = 0u; z + params.seqLength < params.stateWindow; z = z + 1u) {
+    for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+      present_state[z * params.stateSlotStride + ((batch_idx * params.kvNumHeads + head_idx) * HEAD_DIM_K + d) * head_dim_v + dv_idx] = {{ write_scalar("0.0", stateDtype) }};
+    }
+  }
+{% endif %}
+  for (var t = 0u; t < params.seqLength; t = t + 1u) {
+    let bt = batch_idx * params.seqLength + t;
+    let key_base = bt * params.kPackedDim + key_head_idx * HEAD_DIM_K;
+    for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+      key_local[d] = {{ read_scalar("key", "key_base + d", keyDtype) }};
+    }
+{% if gatedDeltaRule %}
+    if (params.decayPackedDim == params.kvNumHeads) {
+      let decay_factor = exp({{ read_scalar("decay", "bt * params.kvNumHeads + head_idx", decayDtype) }});
+      for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+        state[d] = state[d] * decay_factor;
+      }
+    } else {
+      let decay_base = bt * params.decayPackedDim + head_idx * HEAD_DIM_K;
+      for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+        state[d] = state[d] * exp({{ read_scalar("decay", "decay_base + d", decayDtype) }});
+      }
+    }
+    let q_head_0 = (head_idx * params.qNumHeads) / params.kvNumHeads;
+    let out_head_0 = head_idx * heads_per_group;
+    let q0_base = bt * params.qPackedDim + q_head_0 * HEAD_DIM_K;
+    var retrieved = 0.0;
+    var preout = 0.0;
+    var kq_dot = 0.0;
+    for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+      let q_val = {{ read_scalar("query", "q0_base + d", queryDtype) }};
+      retrieved = retrieved + state[d] * key_local[d];
+      preout = preout + state[d] * q_val;
+      kq_dot = kq_dot + key_local[d] * q_val;
+    }
+    var beta_idx = bt * params.kvNumHeads + head_idx;
+    if (params.betaPackedDim == 1u) {
+      beta_idx = bt;
+    }
+    let beta_val = {{ read_scalar("beta", "beta_idx", betaDtype) }};
+    let v_idx = bt * params.vPackedDim + head_idx * head_dim_v + dv_idx;
+    let delta = beta_val * ({{ read_scalar("value", "v_idx", valueDtype) }} - retrieved);
+    let out_idx_0 = bt * packed_out + out_head_0 * head_dim_v + dv_idx;
+    output[out_idx_0] = {{ write_scalar("(preout + delta * kq_dot) * scale", outputDtype) }};
+    for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+      state[d] = state[d] + key_local[d] * delta;
+    }
+{{ emit_serial_query_groups(1) }}
+{% else %}
+    let v_idx = bt * params.vPackedDim + head_idx * head_dim_v + dv_idx;
+    let v_val = {{ read_scalar("value", "v_idx", valueDtype) }};
+    for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+      state[d] = state[d] + key_local[d] * v_val;
+    }
+{{ emit_serial_query_groups(0) }}
+{% endif %}
+{% if hasStateWindow %}
+    // Slot j holds the state after token (seqLength - stateWindow + j), so this
+    // token owns slot (t + stateWindow - seqLength) when that lands in the window.
+    if (t + params.stateWindow >= params.seqLength) {
+      let win_slot = t + params.stateWindow - params.seqLength;
+      for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+        present_state[win_slot * params.stateSlotStride + ((batch_idx * params.kvNumHeads + head_idx) * HEAD_DIM_K + d) * head_dim_v + dv_idx] = {{ write_scalar("state[d]", stateDtype) }};
+      }
+    }
+{% endif %}
+  }
+{% if not hasStateWindow %}
+  for (var d = 0u; d < HEAD_DIM_K; d = d + 1u) {
+    let state_idx = ((batch_idx * params.kvNumHeads + head_idx) * HEAD_DIM_K + d) * head_dim_v + dv_idx;
+    present_state[state_idx] = {{ write_scalar("state[d]", stateDtype) }};
+  }
+{% endif %}
+}

build/webgpu/linear-attention.vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,349 @@

+{% macro read_scalar(name, index, dtype) %}
+{% if dtype == "float16" %}
+f32({{ name }}[{{ index }}]){% else %}{{ name }}[{{ index }}]{% endif %}
+{% endmacro %}
+{% macro write_scalar(expr, dtype) %}
+{% if dtype == "float16" %}
+f16({{ expr }}){% else %}{{ expr }}{% endif %}
+{% endmacro -%}
+{% macro emit_tiled_setup(dvGroups=1) %}
+  let head_dim_k = params.qPackedDim / params.qNumHeads;
+  let head_dim_v = params.vPackedDim / params.kvNumHeads;
+  let n_key_heads = params.kPackedDim / head_dim_k;
+  let heads_per_group = max(1u, params.qNumHeads / params.kvNumHeads);
+  let kv_per_key_head = params.kvNumHeads / n_key_heads;
+{% if dvGroups == 1 %}
+  let dv_tiles = (head_dim_v + TILE_V - 1u) / TILE_V;
+{% else %}
+  // Tile slots per workgroup-index step: each step covers DV_GROUPS value tiles.
+  let dv_tiles = ((head_dim_v + TILE_V - 1u) / TILE_V + DV_GROUPS - 1u) / DV_GROUPS;
+{% endif %}
+  let scale = select(inverseSqrt(f32(head_dim_k)), params.scale, params.scale != 0.0);
+  // 2D-folded flat (batch*head*dv_tile) index: wg.y carries the high bits past
+  // the maxComputeWorkgroupsPerDimension dispatch limit. Reduces to wg.x when nwg.y == 1; the batch_idx >=
+  // params.batchSize guard drops the over-dispatched tail.
+  let workgroup_idx = wg.x + wg.y * nwg.x;
+  let dv_tile_idx = workgroup_idx % dv_tiles;
+  let bh = workgroup_idx / dv_tiles;
+  let head_idx = bh % params.kvNumHeads;
+  let batch_idx = bh / params.kvNumHeads;
+  if (batch_idx >= params.batchSize) {
+    return;
+  }
+{% if dvGroups == 1 %}
+  let dv_start = dv_tile_idx * TILE_V;
+{% else %}
+  let dv_start = (dv_tile_idx * DV_GROUPS + dv_group) * TILE_V;
+{% endif %}
+  let packed_out = max(params.qNumHeads, params.kvNumHeads) * head_dim_v;
+  let key_head_idx = head_idx / kv_per_key_head;
+{%- endmacro -%}
+{% macro emit_vec4_query_groups(first_group) %}
+    for (var qg = {{ first_group }}u; qg < heads_per_group; qg = qg + 1u) {
+      let q_head = (head_idx * params.qNumHeads) / params.kvNumHeads + qg;
+      let out_head = head_idx * heads_per_group + qg;
+      var q_vec = vec4<f32>(0.0);
+      if (lane_active) {
+        let q_base = (bt * params.qPackedDim + q_head * head_dim_k) / VEC4_LANES + lane;
+        q_vec = vec4<f32>(query[q_base]);
+      }
+      var pre_qg: array<f32, TILE_V>;
+{% if useSubgroups %}
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        pre_qg[j] = subgroupAdd(dot(state[j], q_vec));
+      }
+{% else %}
+{% for v in range(redVecs) %}
+      wg_fold[{{ v }}u * WG + tid] = {{ redT }}({% for c in range(redWidth) %}dot(state[{{ v * redWidth + c }}], q_vec){{ ", " if not loop.last else "" }}{% endfor %});
+{% endfor %}
+      workgroupBarrier();
+      // One lane folds each staged slot. LANES is the reduction width, which small
+      // head_dim_k can drive below the slot count, so lanes stride over the slots.
+      for (var sl = lane; sl < {{ redVecs }}u; sl = sl + LANES) {
+        var t = {{ redT }}(0.0);
+        for (var i = 0u; i < LANES; i = i + 1u) {
+          t = t + wg_fold[sl * WG + dv_group * LANES + i];
+        }
+        wg_fold_out[dv_group * {{ redSlots }}u + sl] = t;
+      }
+      workgroupBarrier();
+{% for j in range(tileV) %}
+      pre_qg[{{ j }}] = wg_fold_out[dv_group * {{ redSlots }}u + {{ (j / redWidth)|int }}u]{{ ("." ~ redComps[j % redWidth]) if redWidth > 1 else "" }};
+{% endfor %}
+{% endif %}
+      if (lane == 0u) {
+        let out_base_qg = bt * packed_out + out_head * head_dim_v + dv_start;
+        for (var j = 0u; j < TILE_V; j = j + 1u) {
+          if (dv_start + j < head_dim_v) {
+            output[out_base_qg + j] = {{ write_scalar("pre_qg[j] * scale", queryDtype) }};
+          }
+        }
+      }
+    }
+{%- endmacro %}
+{% set usesDecay = updateRule == "gated" or updateRule == "gated_delta" %}
+{% set redWidth = 4 if tileV % 4 == 0 else (2 if tileV % 2 == 0 else 1) %}
+{% set redVecs = (tileV / redWidth)|int %}
+{% set redT = ("vec" ~ redWidth ~ "<f32>") if redWidth > 1 else "f32" %}
+{% set redComps = ["x", "y", "z", "w"] %}
+{% set usesBeta = updateRule == "delta" or updateRule == "gated_delta" %}
+{% set redSlots = (2 * redVecs + 1) if usesBeta else redVecs %}
+{% if queryDtype == "float16" or stateDtype == "float16" %}
+enable f16;
+{% endif %}
+{% if useSubgroups %}
+enable subgroups;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// LANES threads cooperate on one value tile's reduction axis; DV_GROUPS such groups
+// share a workgroup so the per-token key/query/decay stream is fetched once for all of
+// them instead of once per value tile.
+const LANES: u32 = {{ vec4Lanes }}u;
+const DV_GROUPS: u32 = {{ dvGroups }}u;
+const WG: u32 = LANES * DV_GROUPS;
+const TILE_V: u32 = {{ tileV }}u;
+const VEC4_LANES: u32 = 4u;
+// query and key are bound as vec4, so a lane's four reduction components arrive in one
+// 16-byte fetch instead of four scalar ones. Every row base this shader indexes is a
+// multiple of four -- head_dim_k % 4 == 0 gates the family, and the packed dims are
+// whole numbers of heads -- so the vec4 index is the row base in vec4 units plus the
+// lane. The remaining scalar reads (decay's per-head arm, value, beta) are not aligned
+// groups of four and keep their scalar view.
+{% if not useSubgroups %}
+// No-subgroup tier: shared-memory linear folds replace subgroupAdd. Every quantity a
+// token reduces is staged before one barrier and read back after it, so the barrier
+// count is a property of the token and not of how many quantities it reduces.
+// A lane's TILE_V components are adjacent, so they travel as one {{ redT }} word and
+// fold with {{ redWidth }}-wide adds -- the same per-component summation order as a
+// scalar fold, at a {{ redWidth }}th of the shared-memory transactions.
+// Folds stay within the caller's own lane group: groups own disjoint value tiles and
+// must not see each other's partials.
+var<workgroup> wg_fold: array<{{ redT }}, WG * {{ redSlots }}u>;
+var<workgroup> wg_fold_out: array<{{ redT }}, DV_GROUPS * {{ redSlots }}u>;
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>,
+) {
+  let tid = lid.x;
+  let lane = tid % LANES;
+{% if dvGroups > 1 or not useSubgroups %}
+  let dv_group = tid / LANES;
+{% endif %}
+  let dk_base = lane * VEC4_LANES;
+{{ emit_tiled_setup(dvGroups=dvGroups) }}
+  let lane_active = dk_base < head_dim_k;
+  // state[j] holds 4 consecutive dk rows (the 4 components) for dv slot j.
+  var state: array<vec4<f32>, TILE_V>;
+  for (var j = 0u; j < TILE_V; j = j + 1u) {
+    state[j] = vec4<f32>(0.0);
+  }
+{% if hasPastState %}
+  if (lane_active) {
+{% if hasStateWindow %}
+    // A windowed past_state is read only from slot stateWindow-1, the state after
+    // the last token of the previous call.
+{% endif %}
+    let state_row_base = {% if hasStateWindow %}(params.stateWindow - 1u) * params.stateSlotStride + {% endif %}((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + dk_base) * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        state[j] = vec4<f32>(
+          {{ read_scalar("past_state", "state_row_base + 0u * head_dim_v + j", stateDtype) }},
+          {{ read_scalar("past_state", "state_row_base + 1u * head_dim_v + j", stateDtype) }},
+          {{ read_scalar("past_state", "state_row_base + 2u * head_dim_v + j", stateDtype) }},
+          {{ read_scalar("past_state", "state_row_base + 3u * head_dim_v + j", stateDtype) }},
+        );
+      }
+    }
+  }
+{% endif %}
+{% if hasStateWindow %}
+  // Slots below max(0, stateWindow - seqLength) hold no token from this call.
+  if (lane_active) {
+    let zero_row_base = ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + dk_base) * head_dim_v + dv_start;
+    for (var z = 0u; z + params.seqLength < params.stateWindow; z = z + 1u) {
+      let zb = z * params.stateSlotStride + zero_row_base;
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        if (dv_start + j < head_dim_v) {
+          present_state[zb + 0u * head_dim_v + j] = {{ write_scalar("0.0", stateDtype) }};
+          present_state[zb + 1u * head_dim_v + j] = {{ write_scalar("0.0", stateDtype) }};
+          present_state[zb + 2u * head_dim_v + j] = {{ write_scalar("0.0", stateDtype) }};
+          present_state[zb + 3u * head_dim_v + j] = {{ write_scalar("0.0", stateDtype) }};
+        }
+      }
+    }
+  }
+{% endif %}
+  // A token's key and query do not depend on the recurrent state, so they are fetched one
+  // iteration ahead: the fetch for t+1 is issued before the reductions for t, and its memory
+  // latency overlaps the dependent chain instead of stalling in front of it. The state
+  // recurrence is what serializes this loop, and it leaves the load unit idle otherwise.
+{% macro load_key(token) %}
+    if (lane_active) {
+      let k_base = ({{ token }} * params.kPackedDim + key_head_idx * head_dim_k) / VEC4_LANES + lane;
+      k_next = vec4<f32>(key[k_base]);
+    }
+{%- endmacro %}
+{% if usesBeta %}
+{% macro load_query0(token) %}
+    if (lane_active) {
+      let q0_base = ({{ token }} * params.qPackedDim + q_head_0 * head_dim_k) / VEC4_LANES + lane;
+      q0_next = vec4<f32>(query[q0_base]);
+    }
+{%- endmacro %}
+  let q_head_0 = (head_idx * params.qNumHeads) / params.kvNumHeads;
+  let out_head_0 = head_idx * heads_per_group;
+  var q0_next = vec4<f32>(0.0);
+{% endif %}
+  var k_next = vec4<f32>(0.0);
+  let bt_first = batch_idx * params.seqLength;
+{{ load_key("bt_first") }}
+{% if usesBeta %}
+{{ load_query0("bt_first") }}
+{% endif %}
+  for (var t = 0u; t < params.seqLength; t = t + 1u) {
+    let bt = batch_idx * params.seqLength + t;
+    let k_vec = k_next;
+{% if usesBeta %}
+    let q0_vec = q0_next;
+{% endif %}
+    // Issued here, consumed on the next trip.
+    let bt_next = bt + 1u;
+    if (t + 1u < params.seqLength) {
+{{ load_key("bt_next") }}
+{% if usesBeta %}
+{{ load_query0("bt_next") }}
+{% endif %}
+    }
+{% if usesDecay %}
+    var decay_vec = vec4<f32>(1.0);
+    if (params.decayPackedDim == params.kvNumHeads) {
+      let factor = exp({{ read_scalar("decay", "bt * params.kvNumHeads + head_idx", queryDtype) }});
+      decay_vec = vec4<f32>(factor);
+    } else if (lane_active) {
+      let decay_base = bt * params.decayPackedDim + head_idx * head_dim_k + dk_base;
+      decay_vec = vec4<f32>(
+        exp({{ read_scalar("decay", "decay_base + 0u", queryDtype) }}),
+        exp({{ read_scalar("decay", "decay_base + 1u", queryDtype) }}),
+        exp({{ read_scalar("decay", "decay_base + 2u", queryDtype) }}),
+        exp({{ read_scalar("decay", "decay_base + 3u", queryDtype) }}),
+      );
+    }
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      state[j] = state[j] * decay_vec;
+    }
+{% endif %}
+{% if usesBeta %}
+    var ret_vals: array<f32, TILE_V>;
+    var pre_vals: array<f32, TILE_V>;
+{% if useSubgroups %}
+    let kq_dot = subgroupAdd(dot(k_vec, q0_vec));
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      ret_vals[j] = subgroupAdd(dot(state[j], k_vec));
+      pre_vals[j] = subgroupAdd(dot(state[j], q0_vec));
+    }
+{% else %}
+    // kq and both state projections are known before any of them is consumed, so all
+    // three stage together and cross one barrier pair rather than two.
+{% for v in range(redVecs) %}
+    wg_fold[{{ v }}u * WG + tid] = {{ redT }}({% for c in range(redWidth) %}dot(state[{{ v * redWidth + c }}], k_vec){{ ", " if not loop.last else "" }}{% endfor %});
+    wg_fold[{{ redVecs + v }}u * WG + tid] = {{ redT }}({% for c in range(redWidth) %}dot(state[{{ v * redWidth + c }}], q0_vec){{ ", " if not loop.last else "" }}{% endfor %});
+{% endfor %}
+    wg_fold[{{ 2 * redVecs }}u * WG + tid] = {{ redT }}(dot(k_vec, q0_vec){% for c in range(redWidth - 1) %}, 0.0{% endfor %});
+    workgroupBarrier();
+    for (var sl = lane; sl < {{ redSlots }}u; sl = sl + LANES) {
+      var t = {{ redT }}(0.0);
+      for (var i = 0u; i < LANES; i = i + 1u) {
+        t = t + wg_fold[sl * WG + dv_group * LANES + i];
+      }
+      wg_fold_out[dv_group * {{ redSlots }}u + sl] = t;
+    }
+    workgroupBarrier();
+    let kq_dot = wg_fold_out[dv_group * {{ redSlots }}u + {{ 2 * redVecs }}u]{{ ("." ~ redComps[0]) if redWidth > 1 else "" }};
+{% for j in range(tileV) %}
+    ret_vals[{{ j }}] = wg_fold_out[dv_group * {{ redSlots }}u + {{ (j / redWidth)|int }}u]{{ ("." ~ redComps[j % redWidth]) if redWidth > 1 else "" }};
+    pre_vals[{{ j }}] = wg_fold_out[dv_group * {{ redSlots }}u + {{ redVecs + (j / redWidth)|int }}u]{{ ("." ~ redComps[j % redWidth]) if redWidth > 1 else "" }};
+{% endfor %}
+{% endif %}
+    var beta_idx = bt * params.kvNumHeads + head_idx;
+    if (params.betaPackedDim == 1u) {
+      beta_idx = bt;
+    }
+    let beta_val = {{ read_scalar("beta", "beta_idx", queryDtype) }};
+    let v_base = bt * params.vPackedDim + head_idx * head_dim_v + dv_start;
+    let out_base = bt * packed_out + out_head_0 * head_dim_v + dv_start;
+    var deltas: array<f32, TILE_V>;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        let v_val = {{ read_scalar("value", "v_base + j", queryDtype) }};
+        let delta_j = beta_val * (v_val - ret_vals[j]);
+        deltas[j] = delta_j;
+        if (lane == 0u) {
+          output[out_base + j] = {{ write_scalar("(pre_vals[j] + delta_j * kq_dot) * scale", queryDtype) }};
+        }
+      } else {
+        deltas[j] = 0.0;
+      }
+    }
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      state[j] = state[j] + k_vec * deltas[j];
+    }
+{{ emit_vec4_query_groups(1) }}
+{% else %}
+    let v_base = bt * params.vPackedDim + head_idx * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        let v_val = {{ read_scalar("value", "v_base + j", queryDtype) }};
+        state[j] = state[j] + k_vec * v_val;
+      }
+    }
+{{ emit_vec4_query_groups(0) }}
+{% endif %}
+{% if hasStateWindow %}
+    // Slot j holds the state after token (seqLength - stateWindow + j), so this
+    // token owns slot (t + stateWindow - seqLength) when that lands in the window.
+    if (t + params.stateWindow >= params.seqLength && lane_active) {
+      let win_base = (t + params.stateWindow - params.seqLength) * params.stateSlotStride
+        + ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + dk_base) * head_dim_v + dv_start;
+      for (var j = 0u; j < TILE_V; j = j + 1u) {
+        if (dv_start + j < head_dim_v) {
+          present_state[win_base + 0u * head_dim_v + j] = {{ write_scalar("state[j].x", stateDtype) }};
+          present_state[win_base + 1u * head_dim_v + j] = {{ write_scalar("state[j].y", stateDtype) }};
+          present_state[win_base + 2u * head_dim_v + j] = {{ write_scalar("state[j].z", stateDtype) }};
+          present_state[win_base + 3u * head_dim_v + j] = {{ write_scalar("state[j].w", stateDtype) }};
+        }
+      }
+    }
+{% endif %}
+  }
+{% if not hasStateWindow %}
+  if (lane_active) {
+    let state_row_base = ((batch_idx * params.kvNumHeads + head_idx) * head_dim_k + dk_base) * head_dim_v + dv_start;
+    for (var j = 0u; j < TILE_V; j = j + 1u) {
+      if (dv_start + j < head_dim_v) {
+        present_state[state_row_base + 0u * head_dim_v + j] = {{ write_scalar("state[j].x", stateDtype) }};
+        present_state[state_row_base + 1u * head_dim_v + j] = {{ write_scalar("state[j].y", stateDtype) }};
+        present_state[state_row_base + 2u * head_dim_v + j] = {{ write_scalar("state[j].z", stateDtype) }};
+        present_state[state_row_base + 3u * head_dim_v + j] = {{ write_scalar("state[j].w", stateDtype) }};
+      }
+    }
+  }
+{% endif %}
+}

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

+{
+  "name": "com.microsoft.LinearAttention",
+  "id": "_com_microsoft_linearattention_webgpu_dc21501",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "UnbCbJjzzX17ihI6vVaaXbbgxyNRzghqQ9jFtny3sZs=",
+      "chunk-out.wgsl.jinja": "n2BVpE+dJE8bWsb3Vv9EZalVcE6Bpr0+RYnUxCliSSQ=",
+      "chunk-prep.wgsl.jinja": "U9V7Cld2Zou0sX77ohI+0hmq+UZUsMQohiWrCLFrFm0=",
+      "chunk-scan.wgsl.jinja": "4yrz7CYz19GJAxJ5cBmc4XJt6aXa1ICUi+fg4ma2+mM=",
+      "chunk-ut.wgsl.jinja": "9KqJJCt5PwygvCrBjjy8oglG1ly4HiumH2Im/4Zlz5M=",
+      "linear-attention.scalar.wgsl.jinja": "mc7lY/Xgl1LkvRxiCeaSZ2DwldVlIeoTti2QFqhvCiI=",
+      "linear-attention.serial.wgsl.jinja": "p2AF3fRzazYsMPvGijvV3UyJhxvVxdZJyOhCYLRCJgI=",
+      "linear-attention.vec4.wgsl.jinja": "heG1tNzKu8g1M6PqKtDFED/CzGMDA/qWvlTuZ4Z557M=",
+      "manifest.json": "75pbsCWWrC10tm4NEaVyyoM7nDVgCqxwvj8Bh568M2g=",
+      "test.json": "NtZMMzB3k6UYTuxbvHod7FJlq+HqMV0oYuZi/qhW8sM="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.LinearAttention" }
+}

build/webgpu/test.json ADDED Viewed

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