sync 2e7068faf55e

Browse files

Files changed (7) hide show

README.md +99 -0
build/webgpu/bench.json +293 -0
build/webgpu/manifest.json +535 -0
build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja +117 -0
build/webgpu/metadata.json +19 -0
build/webgpu/qkv-projection.wgsl.jinja +221 -0
build/webgpu/test.json +663 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.MatMulNBitsQkv
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Fuses RMS normalization with three block-quantized attention projections: `A_norm = RMSNorm(A + skip, norm_scale)` (or without `skip`), followed by Q, K, and V projections. The optional fourth output returns `A + skip`. Only 4-bit weights with `block_size = 32` are supported; projection biases, bfloat16, and non-default `accuracy_level` values are not implemented.
+See the [ONNX Runtime `MatMulNBitsQkv` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MatMulNBitsQkv) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `A` | `aT` | `T1` | — | — | Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis. | required |
+| `skip` | `skipT` | `T1` | — | — | Residual added to A before the normalization, with A's shape. | optional |
+| `norm_scale` | `normScaleT` | `T1` | `1` | — | Simplified-layer-normalization (RMS) gain of shape `[K]`. | required |
+| `q_B` | `qBT` | `T2` | `3` | — | Bit-packed uint8 Q weights of shape `(Nq, k_blocks, blob_size)`. | required |
+| `q_scales` | `qScalesT` | `T1` | `2` | — | Per-block Q scales of shape `(Nq, k_blocks)`. Quantization is symmetric: there is no zero-point input, so codes are offset by the midpoint `2^(bits - 1)`. | required |
+| `k_B` | `kBT` | `T2` | `3` | — | Bit-packed K weights of shape `(Nkv, k_blocks, blob_size)`. | required |
+| `k_scales` | `kScalesT` | `T1` | `2` | — | Per-block K scales of shape `(Nkv, k_blocks)`. | required |
+| `v_B` | `vBT` | `T2` | `3` | — | Bit-packed V weights of shape `(Nkv, k_blocks, blob_size)`. | required |
+| `v_scales` | `vScalesT` | `T1` | `2` | — | Per-block V scales of shape `(Nkv, k_blocks)`. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Q` | `qT` | `T1` | same as `A` | derived; see description | Query projection: A's leading axes with a trailing Nq. | required |
+| `K` | `kT` | `T1` | same as `A` | derived; see description | Key projection: A's leading axes with a trailing Nkv. | required |
+| `V` | `vT` | `T1` | same as `A` | derived; see description | Value projection: A's leading axes with a trailing Nkv. | required |
+| `input_skip_bias_sum` | `residualT` | `T1` | same as `A` | same as `A` | The residual sum A + skip, with A's shape. Requires the skip input. | optional |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `accuracy_level` | `0` | Minimum internal accuracy level, following MatMulNBits semantics; this implementation supports the standard default 0. |
+| `bits` | `4` | Bit width used to quantize all three weight matrices; only 4 is supported. |
+| `epsilon` | `9.999999974752427e-7` | Epsilon used by the simplified layer-normalization reduction. |
+| `K` | — | Input feature dimension shared by the normalized input and all projection weights. |
+| `Nq` | — | Output feature dimension of the Q projection. |
+| `Nkv` | — | Output feature dimension shared by the K and V projections. |
+| `block_size` | — | Size of each quantization block along K; only 32 is supported. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T1` | `float32`, `float16` |
+| `T2` | `uint8` |
+## 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
+- [`matmul-nbits-fused-rms-norm.wgsl.jinja`](build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja)
+- [`qkv-projection.wgsl.jinja`](build/webgpu/qkv-projection.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.MatMulNBitsQkv", { version: 1 });
+const { qT, kT, vT } = await kernel({
+  aT: { data: aTData, shape: [2, 32] },
+  normScaleT: { data: normScaleTData, shape: [32] },
+  qBT: { data: qBTData, shape: [5, 1, 16] },
+  qScalesT: { data: qScalesTData, shape: [5, 1] },
+  kBT: { data: kBTData, shape: [3, 1, 16] },
+  kScalesT: { data: kScalesTData, shape: [3, 1] },
+  vBT: { data: vBTData, shape: [3, 1, 16] },
+  vScalesT: { data: vScalesTData, shape: [3, 1] },
+}, {
+  attrs: {
+    K: 32,
+    Nq: 5,
+    Nkv: 3,
+    block_size: 32,
+  },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,293 @@

+{
+  "op": "com.microsoft.MatMulNBitsQkv",
+  "tunableSpace": { "TILE_N": [4, 8, 16], "LANES": [4, 8, 16] },
+  "cases": [
+    {
+      "name": "qkv-q4-decode-k2048-nq2048-nkv512",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "attrs": { "K": 2048, "Nq": 2048, "Nkv": 512, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 9101, "scale": 1 },
+        "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 9102, "scale": 1 },
+        "qBT": { "shape": [2048, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 9103, "scale": 255 },
+        "qScalesT": { "shape": [2048, 64], "dtype": "float32", "dist": "normal", "seed": 9104, "scale": 0.05 },
+        "kBT": { "shape": [512, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 9105, "scale": 255 },
+        "kScalesT": { "shape": [512, 64], "dtype": "float32", "dist": "normal", "seed": 9106, "scale": 0.05 },
+        "vBT": { "shape": [512, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 9107, "scale": 255 },
+        "vScalesT": { "shape": [512, 64], "dtype": "float32", "dist": "normal", "seed": 9108, "scale": 0.05 }
+      },
+      "outputs": {
+        "qT": { "shape": [1, 2048], "dtype": "float32" },
+        "kT": { "shape": [1, 512], "dtype": "float32" },
+        "vT": { "shape": [1, 512], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(2048 + 512 + 512) * 64 * 16 * 4" }] }
+    },
+    {
+      "name": "qkv-q4-llama-decode-k4096-nq4096-nkv4096",
+      "preset": "model",
+      "provenance": {
+        "notes": "Llama class defaults (hidden_size 4096, num_attention_heads 32, no GQA so num_key_value_heads matches) at a decode step. Bytes are GPU storage, not the packed blob: this backend gives a uint8 tensor one byte per u32 element, so a 4-bit code costs 2 bytes of traffic and the projection moves four times its on-disk size. Counting the blob understated the traffic 4x and reported this kernel at a tenth of the bandwidth rail when it runs at about half."
+      },
+      "vars": { "dtype": "float32", "tokens": 1, "K": 4096, "Nq": 4096, "Nkv": 4096 },
+      "attrs": { "K": 4096, "Nq": 4096, "Nkv": 4096, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 7100, "scale": 0.5 },
+        "normScaleT": {
+          "shape": [4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7101,
+          "scale": 0.1,
+          "offset": 1
+        },
+        "qBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7102, "min": 0, "max": 255 },
+        "qScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7103,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "kBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7104, "min": 0, "max": 255 },
+        "kScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7105,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "vBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7106, "min": 0, "max": 255 },
+        "vScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7107,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": {
+        "qT": { "shape": [1, 4096], "dtype": "float32" },
+        "kT": { "shape": [1, 4096], "dtype": "float32" },
+        "vT": { "shape": [1, 4096], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(args.Nq + 2 * args.Nkv) * args.K * 2" }] }
+    },
+    {
+      "name": "qkv-q4-mistral-decode-k4096-nq4096-nkv1024",
+      "preset": "model",
+      "provenance": {
+        "notes": "Mistral class defaults (hidden_size 4096, 32 query heads, 8 key/value heads, head_dim 128) -- the GQA case, where K and V projections are a quarter of Q. Bytes are GPU storage, not the packed blob: this backend gives a uint8 tensor one byte per u32 element, so a 4-bit code costs 2 bytes of traffic and the projection moves four times its on-disk size. Counting the blob understated the traffic 4x and reported this kernel at a tenth of the bandwidth rail when it runs at about half."
+      },
+      "vars": { "dtype": "float32", "tokens": 1, "K": 4096, "Nq": 4096, "Nkv": 1024 },
+      "attrs": { "K": 4096, "Nq": 4096, "Nkv": 1024, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 7200, "scale": 0.5 },
+        "normScaleT": {
+          "shape": [4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7201,
+          "scale": 0.1,
+          "offset": 1
+        },
+        "qBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7202, "min": 0, "max": 255 },
+        "qScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7203,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "kBT": { "shape": [1024, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7204, "min": 0, "max": 255 },
+        "kScalesT": {
+          "shape": [1024, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7205,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "vBT": { "shape": [1024, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7206, "min": 0, "max": 255 },
+        "vScalesT": {
+          "shape": [1024, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7207,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": {
+        "qT": { "shape": [1, 4096], "dtype": "float32" },
+        "kT": { "shape": [1, 1024], "dtype": "float32" },
+        "vT": { "shape": [1, 1024], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(args.Nq + 2 * args.Nkv) * args.K * 2" }] }
+    },
+    {
+      "name": "qkv-q4-gemma3-decode-k2304-nq2048-nkv1024",
+      "preset": "model",
+      "provenance": {
+        "notes": "Gemma3 class defaults (hidden_size 2304, 8 query heads, 4 key/value heads, head_dim 256); the projection is wider than hidden, so K does not equal Nq. Bytes are GPU storage, not the packed blob: this backend gives a uint8 tensor one byte per u32 element, so a 4-bit code costs 2 bytes of traffic and the projection moves four times its on-disk size. Counting the blob understated the traffic 4x and reported this kernel at a tenth of the bandwidth rail when it runs at about half."
+      },
+      "vars": { "dtype": "float32", "tokens": 1, "K": 2304, "Nq": 2048, "Nkv": 1024 },
+      "attrs": { "K": 2304, "Nq": 2048, "Nkv": 1024, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 2304], "dtype": "float32", "dist": "normal", "seed": 7300, "scale": 0.5 },
+        "normScaleT": {
+          "shape": [2304],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7301,
+          "scale": 0.1,
+          "offset": 1
+        },
+        "qBT": { "shape": [2048, 72, 16], "dtype": "uint8", "dist": "uniform", "seed": 7302, "min": 0, "max": 255 },
+        "qScalesT": {
+          "shape": [2048, 72],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7303,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "kBT": { "shape": [1024, 72, 16], "dtype": "uint8", "dist": "uniform", "seed": 7304, "min": 0, "max": 255 },
+        "kScalesT": {
+          "shape": [1024, 72],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7305,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "vBT": { "shape": [1024, 72, 16], "dtype": "uint8", "dist": "uniform", "seed": 7306, "min": 0, "max": 255 },
+        "vScalesT": {
+          "shape": [1024, 72],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7307,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": {
+        "qT": { "shape": [1, 2048], "dtype": "float32" },
+        "kT": { "shape": [1, 1024], "dtype": "float32" },
+        "vT": { "shape": [1, 1024], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(args.Nq + 2 * args.Nkv) * args.K * 2" }] }
+    },
+    {
+      "name": "qkv-q4-phi3-decode-k3072-nq3072-nkv3072",
+      "preset": "model",
+      "provenance": {
+        "notes": "Phi-3 class defaults (hidden_size 3072, num_attention_heads 32, no GQA). Bytes are GPU storage, not the packed blob: this backend gives a uint8 tensor one byte per u32 element, so a 4-bit code costs 2 bytes of traffic and the projection moves four times its on-disk size. Counting the blob understated the traffic 4x and reported this kernel at a tenth of the bandwidth rail when it runs at about half."
+      },
+      "vars": { "dtype": "float32", "tokens": 1, "K": 3072, "Nq": 3072, "Nkv": 3072 },
+      "attrs": { "K": 3072, "Nq": 3072, "Nkv": 3072, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 3072], "dtype": "float32", "dist": "normal", "seed": 7400, "scale": 0.5 },
+        "normScaleT": {
+          "shape": [3072],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7401,
+          "scale": 0.1,
+          "offset": 1
+        },
+        "qBT": { "shape": [3072, 96, 16], "dtype": "uint8", "dist": "uniform", "seed": 7402, "min": 0, "max": 255 },
+        "qScalesT": {
+          "shape": [3072, 96],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7403,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "kBT": { "shape": [3072, 96, 16], "dtype": "uint8", "dist": "uniform", "seed": 7404, "min": 0, "max": 255 },
+        "kScalesT": {
+          "shape": [3072, 96],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7405,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "vBT": { "shape": [3072, 96, 16], "dtype": "uint8", "dist": "uniform", "seed": 7406, "min": 0, "max": 255 },
+        "vScalesT": {
+          "shape": [3072, 96],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7407,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": {
+        "qT": { "shape": [1, 3072], "dtype": "float32" },
+        "kT": { "shape": [1, 3072], "dtype": "float32" },
+        "vT": { "shape": [1, 3072], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(args.Nq + 2 * args.Nkv) * args.K * 2" }] }
+    },
+    {
+      "name": "qkv-q4-llama-spec8-k4096-nq4096-nkv4096",
+      "preset": "model",
+      "provenance": {
+        "notes": "Llama class defaults with 8 rows, the shape speculative decoding verifies in one pass. Bytes are GPU storage, not the packed blob: this backend gives a uint8 tensor one byte per u32 element, so a 4-bit code costs 2 bytes of traffic and the projection moves four times its on-disk size. Counting the blob understated the traffic 4x and reported this kernel at a tenth of the bandwidth rail when it runs at about half."
+      },
+      "vars": { "dtype": "float32", "tokens": 8, "K": 4096, "Nq": 4096, "Nkv": 4096 },
+      "attrs": { "K": 4096, "Nq": 4096, "Nkv": 4096, "block_size": 32 },
+      "inputs": {
+        "aT": { "shape": [8, 4096], "dtype": "float32", "dist": "normal", "seed": 7500, "scale": 0.5 },
+        "normScaleT": {
+          "shape": [4096],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7501,
+          "scale": 0.1,
+          "offset": 1
+        },
+        "qBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7502, "min": 0, "max": 255 },
+        "qScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7503,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "kBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7504, "min": 0, "max": 255 },
+        "kScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7505,
+          "min": 0.005,
+          "max": 0.05
+        },
+        "vBT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "uniform", "seed": 7506, "min": 0, "max": 255 },
+        "vScalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 7507,
+          "min": 0.005,
+          "max": 0.05
+        }
+      },
+      "outputs": {
+        "qT": { "shape": [8, 4096], "dtype": "float32" },
+        "kT": { "shape": [8, 4096], "dtype": "float32" },
+        "vT": { "shape": [8, 4096], "dtype": "float32" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(args.Nq + 2 * args.Nkv) * args.K * 2" }] }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,535 @@

+{
+  "domain": "com.microsoft",
+  "name": "MatMulNBitsQkv",
+  "sinceVersion": 1,
+  "description": "Fuses RMS normalization with three block-quantized attention projections: `A_norm = RMSNorm(A + skip, norm_scale)` (or without `skip`), followed by Q, K, and V projections. The optional fourth output returns `A + skip`. Only 4-bit weights with `block_size = 32` are supported; projection biases, bfloat16, and non-default `accuracy_level` values are not implemented.",
+  "inputs": [
+    {
+      "role": "A",
+      "dtype": "T1",
+      "description": "Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis."
+    },
+    {
+      "role": "skip",
+      "dtype": "T1",
+      "optional": true,
+      "description": "Residual added to A before the normalization, with A's shape."
+    },
+    {
+      "role": "norm_scale",
+      "dtype": "T1",
+      "rank": 1,
+      "description": "Simplified-layer-normalization (RMS) gain of shape `[K]`."
+    },
+    {
+      "role": "q_B",
+      "dtype": "T2",
+      "rank": 3,
+      "description": "Bit-packed uint8 Q weights of shape `(Nq, k_blocks, blob_size)`."
+    },
+    {
+      "role": "q_scales",
+      "dtype": "T1",
+      "rank": 2,
+      "description": "Per-block Q scales of shape `(Nq, k_blocks)`. Quantization is symmetric: there is no zero-point input, so codes are offset by the midpoint `2^(bits - 1)`."
+    },
+    {
+      "role": "k_B",
+      "dtype": "T2",
+      "rank": 3,
+      "description": "Bit-packed K weights of shape `(Nkv, k_blocks, blob_size)`."
+    },
+    { "role": "k_scales", "dtype": "T1", "rank": 2, "description": "Per-block K scales of shape `(Nkv, k_blocks)`." },
+    {
+      "role": "v_B",
+      "dtype": "T2",
+      "rank": 3,
+      "description": "Bit-packed V weights of shape `(Nkv, k_blocks, blob_size)`."
+    },
+    { "role": "v_scales", "dtype": "T1", "rank": 2, "description": "Per-block V scales of shape `(Nkv, k_blocks)`." }
+  ],
+  "outputs": [
+    {
+      "role": "Q",
+      "dtype": "T1",
+      "rank": "ranks.aT",
+      "shape": "shapes.aT[:-1] + [attrs.Nq]",
+      "description": "Query projection: A's leading axes with a trailing Nq."
+    },
+    {
+      "role": "K",
+      "dtype": "T1",
+      "rank": "ranks.aT",
+      "shape": "shapes.aT[:-1] + [attrs.Nkv]",
+      "description": "Key projection: A's leading axes with a trailing Nkv."
+    },
+    {
+      "role": "V",
+      "dtype": "T1",
+      "rank": "ranks.aT",
+      "shape": "shapes.aT[:-1] + [attrs.Nkv]",
+      "description": "Value projection: A's leading axes with a trailing Nkv."
+    },
+    {
+      "role": "input_skip_bias_sum",
+      "dtype": "T1",
+      "rank": "ranks.aT",
+      "optional": true,
+      "shape": "shapes.aT",
+      "description": "The residual sum A + skip, with A's shape. Requires the skip input."
+    }
+  ],
+  "attributes": { "accuracy_level": 0, "bits": 4, "epsilon": 9.999999974752427e-7 },
+  "attributeDescriptions": {
+    "K": "Input feature dimension shared by the normalized input and all projection weights.",
+    "Nq": "Output feature dimension of the Q projection.",
+    "Nkv": "Output feature dimension shared by the K and V projections.",
+    "accuracy_level": "Minimum internal accuracy level, following MatMulNBits semantics; this implementation supports the standard default 0.",
+    "bits": "Bit width used to quantize all three weight matrices; only 4 is supported.",
+    "block_size": "Size of each quantization block along K; only 32 is supported.",
+    "epsilon": "Epsilon used by the simplified layer-normalization reduction."
+  },
+  "attributeConstraints": {
+    "K": { "required": true },
+    "Nq": { "required": true },
+    "Nkv": { "required": true },
+    "accuracy_level": { "values": [0] },
+    "bits": { "values": [4] },
+    "block_size": { "required": true, "values": [32] }
+  },
+  "typeConstraints": { "T1": ["float32", "float16"], "T2": ["uint8"] },
+  "args": {
+    "aT": { "kind": "tensor", "semantic": "A", "role": "input" },
+    "skipT": { "kind": "tensor", "semantic": "skip", "role": "input", "required": false },
+    "normScaleT": { "kind": "tensor", "semantic": "norm_scale", "role": "weights" },
+    "qBT": { "kind": "tensor", "semantic": "q_B", "role": "weights" },
+    "qScalesT": { "kind": "tensor", "semantic": "q_scales", "role": "weights" },
+    "kBT": { "kind": "tensor", "semantic": "k_B", "role": "weights" },
+    "kScalesT": { "kind": "tensor", "semantic": "k_scales", "role": "weights" },
+    "vBT": { "kind": "tensor", "semantic": "v_B", "role": "weights" },
+    "vScalesT": { "kind": "tensor", "semantic": "v_scales", "role": "weights" },
+    "qT": { "kind": "tensor", "semantic": "Q", "role": "output" },
+    "kT": { "kind": "tensor", "semantic": "K", "role": "output" },
+    "vT": { "kind": "tensor", "semantic": "V", "role": "output" },
+    "residualT": { "kind": "tensor", "semantic": "input_skip_bias_sum", "role": "output", "required": false }
+  },
+  "tunables": { "TILE_N": 8, "LANES": 8, "NORM_WORKGROUP_SIZE": 128, "ROW_TILE": 8 },
+  "derive": {
+    "aRows": "numel(shapes.aT) / max(1, attrs.K)",
+    "rowTile": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)",
+    "rowGroups": "ceilDiv(aRows, rowTile)",
+    "kBlocks": "dim(shapes.qBT, 1)",
+    "blobSize": "dim(shapes.qBT, 2)",
+    "codesPerByte": "8 / attrs.bits",
+    "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
+    "pairSharesWord": "codesPerByte >= 2",
+    "epsilonValue": "attrs.epsilon",
+    "weightShapeOk": "ranks.qBT == 3 and ranks.kBT == 3 and ranks.vBT == 3 and dim(shapes.qBT, 0) == attrs.Nq and dim(shapes.kBT, 0) == attrs.Nkv and dim(shapes.vBT, 0) == attrs.Nkv and dim(shapes.kBT, 1) == kBlocks and dim(shapes.vBT, 1) == kBlocks and dim(shapes.kBT, 2) == blobSize and dim(shapes.vBT, 2) == blobSize and kBlocks == ceilDiv(attrs.K, attrs.block_size) and blobSize * 8 == attrs.block_size * attrs.bits",
+    "scaleShapeOk": "ranks.qScalesT == 2 and ranks.kScalesT == 2 and ranks.vScalesT == 2 and dim(shapes.qScalesT, 0) == attrs.Nq and dim(shapes.qScalesT, 1) == kBlocks and dim(shapes.kScalesT, 0) == attrs.Nkv and dim(shapes.kScalesT, 1) == kBlocks and dim(shapes.vScalesT, 0) == attrs.Nkv and dim(shapes.vScalesT, 1) == kBlocks",
+    "ioShapeOk": "(ranks.aT == 2 or ranks.aT == 3) and dim(shapes.aT, ranks.aT - 1) == attrs.K and ranks.qT == ranks.aT and ranks.kT == ranks.aT and ranks.vT == ranks.aT and dim(shapes.qT, ranks.qT - 1) == attrs.Nq and dim(shapes.kT, ranks.kT - 1) == attrs.Nkv and dim(shapes.vT, ranks.vT - 1) == attrs.Nkv and sameShape(prefix(shapes.qT, ranks.qT - 1), prefix(shapes.aT, ranks.aT - 1)) and sameShape(prefix(shapes.kT, ranks.kT - 1), prefix(shapes.aT, ranks.aT - 1)) and sameShape(prefix(shapes.vT, ranks.vT - 1), prefix(shapes.aT, ranks.aT - 1))",
+    "dtypeOk": "tensorDtypes.qBT == \"uint8\" and tensorDtypes.kBT == \"uint8\" and tensorDtypes.vBT == \"uint8\" and tensorDtypes.qScalesT == tensorDtypes.aT and tensorDtypes.kScalesT == tensorDtypes.aT and tensorDtypes.vScalesT == tensorDtypes.aT and tensorDtypes.qT == tensorDtypes.aT and tensorDtypes.kT == tensorDtypes.aT and tensorDtypes.vT == tensorDtypes.aT and tensorDtypes.normScaleT == tensorDtypes.aT and f16Ok(tensorDtypes.aT)",
+    "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)",
+    "normContractOk": "ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)",
+    "qkvShapeOk": "weightShapeOk and scaleShapeOk and ioShapeOk and dtypeOk and lanesPow2 and normContractOk and pairSharesWord and attrs.K > 0 and attrs.Nq > 0 and attrs.Nkv > 0",
+    "projectionTiles": "ceilDiv(attrs.Nq, tunables.TILE_N) + 2 * ceilDiv(attrs.Nkv, tunables.TILE_N)",
+    "dispatchFits": "projectionTiles <= device.limits.maxComputeWorkgroupsPerDimension and aRows <= device.limits.maxComputeWorkgroupsPerDimension and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX"
+  },
+  "constants": {
+    "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "usesF16": "tensorDtypes.aT == \"float16\"",
+    "K": "attrs.K",
+    "nq": "attrs.Nq",
+    "nkv": "attrs.Nkv",
+    "blockSize": "attrs.block_size",
+    "kBlocks": "kBlocks",
+    "blobSize": "blobSize",
+    "bits": "attrs.bits",
+    "codesPerByte": "codesPerByte",
+    "codeMask": "codeMask",
+    "defaultZero": "\"8.0\"",
+    "tileN": "tunables.TILE_N",
+    "lanes": "tunables.LANES",
+    "hidden": "attrs.K",
+    "workgroupSize": "tunables.NORM_WORKGROUP_SIZE",
+    "epsilon": "epsilonValue",
+    "hasSkip": "present.skipT",
+    "writeResidual": "present.residualT",
+    "K_LEN": "attrs.K",
+    "rowTile": "rowTile",
+    "rows": "aRows"
+  },
+  "bindingSets": {
+    "normFull": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "skip",
+        "arg": "skipT",
+        "semantic": "skip",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "norm_scale",
+        "arg": "normScaleT",
+        "semantic": "norm_scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar",
+        "length": "$K_LEN"
+      },
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "residual",
+        "arg": "residualT",
+        "semantic": "input_skip_bias_sum",
+        "buffer": { "type": "storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
+      }
+    ],
+    "normSkip": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "skip",
+        "arg": "skipT",
+        "semantic": "skip",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "norm_scale",
+        "arg": "normScaleT",
+        "semantic": "norm_scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar",
+        "length": "$K_LEN"
+      },
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
+      }
+    ],
+    "normOnly": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "norm_scale",
+        "arg": "normScaleT",
+        "semantic": "norm_scale",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar",
+        "length": "$K_LEN"
+      },
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "rows", "type": "u32", "value": "aRows" }] }
+      }
+    ],
+    "projection": [
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "q_b",
+        "arg": "qBT",
+        "semantic": "q_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "q_scales",
+        "arg": "qScalesT",
+        "semantic": "q_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "k_b",
+        "arg": "kBT",
+        "semantic": "k_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "k_scales",
+        "arg": "kScalesT",
+        "semantic": "k_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      {
+        "name": "v_b",
+        "arg": "vBT",
+        "semantic": "v_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "v_scales",
+        "arg": "vScalesT",
+        "semantic": "v_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "q", "arg": "qT", "semantic": "Q", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
+      { "name": "k", "arg": "kT", "semantic": "K", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
+      { "name": "v", "arg": "vT", "semantic": "V", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
+    ],
+    "projectionQ": [
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "q_b",
+        "arg": "qBT",
+        "semantic": "q_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "q_scales",
+        "arg": "qScalesT",
+        "semantic": "q_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "q", "arg": "qT", "semantic": "Q", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
+    ],
+    "projectionK": [
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "k_b",
+        "arg": "kBT",
+        "semantic": "k_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "k_scales",
+        "arg": "kScalesT",
+        "semantic": "k_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "k", "arg": "kT", "semantic": "K", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
+    ],
+    "projectionV": [
+      { "name": "normed", "semantic": "normedA", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "v_b",
+        "arg": "vBT",
+        "semantic": "v_B",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "name": "v_scales",
+        "arg": "vScalesT",
+        "semantic": "v_scales",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "v", "arg": "vT", "semantic": "V", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
+    ]
+  },
+  "variants": [
+    {
+      "id": "norm",
+      "priority": 20,
+      "when": ["qkvShapeOk", "dispatchFits", "not present.skipT", "not present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normOnly",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "main",
+          "name": "MatMulNBitsQkv.Projection",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"\"" } },
+          "bindings": "projection",
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" }
+        }
+      ]
+    },
+    {
+      "id": "split_norm",
+      "priority": 10,
+      "when": ["qkvShapeOk", "dispatchFits", "not present.skipT", "not present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normOnly",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "q",
+          "name": "MatMulNBitsQkv.ProjectionQ",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"q\"" } },
+          "bindings": "projectionQ",
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "k",
+          "name": "MatMulNBitsQkv.ProjectionK",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"k\"" } },
+          "bindings": "projectionK",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "v",
+          "name": "MatMulNBitsQkv.ProjectionV",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"v\"" } },
+          "bindings": "projectionV",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        }
+      ]
+    },
+    {
+      "id": "skip",
+      "priority": 20,
+      "when": ["qkvShapeOk", "dispatchFits", "present.skipT", "not present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normSkip",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "main",
+          "name": "MatMulNBitsQkv.Projection",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"\"" } },
+          "bindings": "projection",
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" }
+        }
+      ]
+    },
+    {
+      "id": "split_skip",
+      "priority": 10,
+      "when": ["qkvShapeOk", "dispatchFits", "present.skipT", "not present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normSkip",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "q",
+          "name": "MatMulNBitsQkv.ProjectionQ",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"q\"" } },
+          "bindings": "projectionQ",
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "k",
+          "name": "MatMulNBitsQkv.ProjectionK",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"k\"" } },
+          "bindings": "projectionK",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "v",
+          "name": "MatMulNBitsQkv.ProjectionV",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"v\"" } },
+          "bindings": "projectionV",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        }
+      ]
+    },
+    {
+      "id": "skipsum",
+      "priority": 20,
+      "when": ["qkvShapeOk", "dispatchFits", "present.skipT", "present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normFull",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "main",
+          "name": "MatMulNBitsQkv.Projection",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"\"" } },
+          "bindings": "projection",
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" }
+        }
+      ]
+    },
+    {
+      "id": "split_skipsum",
+      "priority": 10,
+      "when": ["qkvShapeOk", "dispatchFits", "present.skipT", "present.residualT"],
+      "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
+      "passes": [
+        {
+          "id": "norm",
+          "name": "MatMulNBitsQkv.RmsNorm",
+          "shader": "matmul-nbits-fused-rms-norm.wgsl.jinja",
+          "bindings": "normFull",
+          "dispatch": { "workgroups": "aRows" }
+        },
+        {
+          "id": "q",
+          "name": "MatMulNBitsQkv.ProjectionQ",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"q\"" } },
+          "bindings": "projectionQ",
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "k",
+          "name": "MatMulNBitsQkv.ProjectionK",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"k\"" } },
+          "bindings": "projectionK",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        },
+        {
+          "id": "v",
+          "name": "MatMulNBitsQkv.ProjectionV",
+          "source": { "shader": "qkv-projection.wgsl.jinja", "inputs": { "singleProjection": "\"v\"" } },
+          "bindings": "projectionV",
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tunables.TILE_N)", "y": "rowGroups" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,117 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Shared MatMulNBits MLP/QKV normalization pass.
+//   normed[row, d] = (A + skip)[row, d] * inverseSqrt(mean_d((A + skip)^2) + eps) * norm_scale[d]
+// One workgroup owns one row. Every intermediate stays in f32 and `normed` is an
+// f32 scratch tensor, so the projection pass reads exactly what the fused
+// single-dispatch kernel keeps in registers -- the two paths agree for float16
+// inputs instead of differing by one narrowing.
+const HIDDEN: u32 = {{ hidden }}u;
+const WG: u32 = {{ workgroupSize }}u;
+const EPSILON: f32 = {{ epsilon }};
+var<workgroup> partial: array<f32, WG>;
+{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" %}
+{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
+{%- else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
+{%- endif %}
+{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
+  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
+  loop {
+{% if form == "head" %}
+{% if breakInline %}
+    if ({{ svar }} == 0u) { break; }
+{% else %}
+    if ({{ svar }} == 0u) {
+      break;
+    }
+{% endif %}
+{% endif %}
+{% if bodyInline %}
+    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
+{% else %}
+    if ({{ idx }} < {{ svar }}) {
+{% for a in arrays %}
+      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
+{% endfor %}
+    }
+{% endif %}
+{% if form == "head" %}
+{% if barrierFirst %}
+    workgroupBarrier();
+    {{ svar }} = {{ svar }} / 2u;
+{% else %}
+    {{ svar }} = {{ svar }} / 2u;
+    workgroupBarrier();
+{% endif %}
+{% else %}
+    workgroupBarrier();
+    if ({{ svar }} == 1u) {
+      break;
+    }
+    {{ svar }} = {{ svar }} / 2u;
+{% endif %}
+  }
+{%- endmacro %}
+// Reusing partial after this reduction requires a barrier between the read of
+// partial[0] and the next write, or the next round can race the prior readers.
+{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
+fn reduce_sum(value: f32, tid: u32) -> f32 {
+  partial[tid] = value;
+  workgroupBarrier();
+{{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
+{% if trailingBarrier %}
+  let total = partial[0];
+  workgroupBarrier();
+  return total;
+{% else %}
+  return partial[0];
+{% endif %}
+}
+fn row_value(index: u32) -> f32 {
+{% if hasSkip %}
+  return f32(a[index]) + f32(skip[index]);
+{% else %}
+  return f32(a[index]);
+{% 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>) {
+  // 2D-folded row index: wg.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit. Reduces to wg.x when nwg.y == 1.
+  let row = wg.x + wg.y * nwg.x;
+  if (row >= params.rows) {
+    return;
+  }
+  let tid = lid.x;
+  let base = row * HIDDEN;
+  var local_sq = 0.0;
+  for (var d = tid; d < HIDDEN; d = d + WG) {
+    let value = row_value(base + d);
+    local_sq = local_sq + value * value;
+  }
+  let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(HIDDEN) + EPSILON);
+  for (var d = tid; d < HIDDEN; d = d + WG) {
+    let index = base + d;
+    let value = row_value(index);
+{% if writeResidual %}
+    residual[index] = {{ scalar }}(value);
+{% endif %}
+    normed[index] = value * inv * f32(norm_scale[d]);
+  }
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,19 @@

+{
+  "name": "com.microsoft.MatMulNBitsQkv",
+  "id": "_com_microsoft_matmulnbitsqkv_webgpu_07222b0",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "TMTSE3iJebIWYjPDDVFyPwWC3YdOekG83a3gOGp3Xao=",
+      "manifest.json": "xLf6qGf+IK/6nqf+JAlvzmyVSbm5jUarc5DL/IVW3w0=",
+      "matmul-nbits-fused-rms-norm.wgsl.jinja": "4TI9Mc/RZWVgMl6NVG2iwe5OpvqQmswhe6xtGvctNmI=",
+      "qkv-projection.wgsl.jinja": "sINPOc7UfVRs5YFgICljS7B6QxhlPuL6IVEEatlIaMY=",
+      "test.json": "Ttc9zF84a34u+oxrvXeKL3Wah8leZYmRrkN0Ddw+P0E="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulNBitsQkv" }
+}

build/webgpu/qkv-projection.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,221 @@

+{% macro matmul_nbits_packed_code(fn="packed_weight", buffer="b", kBlocks="params.kBlocks", blobSize="params.blobSize", bits=4) %}
+fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
+{% if bits == 2 %}
+  let byte_index = offset / 4u;
+  let shift = (offset % 4u) * 2u;
+  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
+  return ({{ buffer }}[packed_index] >> shift) & 3u;
+{% elif bits == 4 %}
+  let byte_index = offset / 2u;
+  let shift = (offset % 2u) * 4u;
+  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + byte_index;
+  return ({{ buffer }}[packed_index] >> shift) & 15u;
+{% else %}
+  let packed_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
+  return {{ buffer }}[packed_index] & 255u;
+{% endif %}
+}
+{%- endmacro %}
+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// com.microsoft.MatMulNBitsQkv, projection pass.
+//   Q[row, n] = dot(A_norm[row], q_weight[n]), and likewise K and V.
+{% if source.singleProjection %}
+// This specialization binds and computes one projection.
+{% else %}
+// All three projections read the same normalized row, so one dispatch covers
+// them: the column-tile grid is laid out as Q's tiles, then K's, then V's, and
+// a workgroup's tile index alone decides which projection it serves. That keeps
+// the branch workgroup-uniform -- no lane in a workgroup ever takes a different
+// arm -- while still reading each activation row once per tile instead of once
+// per projection.
+{% endif %}
+// Within a tile, LANES threads cooperate on one output column: thread
+// (column, lane) strides the reduction axis by LANES and the partials are
+// folded inside the column's own lane group. Codes dequantize as
+// (code - ZERO) * block_scale with the scale hoisted out of the inner loop;
+// this operator has no zero-point input, so ZERO is the symmetric midpoint.
+// A workgroup also covers ROW_TILE activation rows, reusing each unpacked code
+// across their accumulators.
+const K: u32 = {{ K }}u;
+{% if source.singleProjection == "q" %}
+const N: u32 = {{ nq }}u;
+{% elif source.singleProjection %}
+const N: u32 = {{ nkv }}u;
+{% else %}
+const NQ: u32 = {{ nq }}u;
+const NKV: u32 = {{ nkv }}u;
+{% endif %}
+const BLOCK_SIZE: u32 = {{ blockSize }}u;
+const KBLOCKS: u32 = {{ kBlocks }}u;
+const BLOB_SIZE: u32 = {{ blobSize }}u;
+const TILE_N: u32 = {{ tileN }}u;
+const LANES: u32 = {{ lanes }}u;
+const ROW_TILE: u32 = {{ rowTile }}u;
+const ROWS: u32 = {{ rows }}u;
+const WG: u32 = TILE_N * LANES;
+const ZERO: f32 = {{ defaultZero }};
+{% if not source.singleProjection %}
+const Q_TILES: u32 = (NQ + TILE_N - 1u) / TILE_N;
+const KV_TILES: u32 = (NKV + TILE_N - 1u) / TILE_N;
+{% endif %}
+const BITS: u32 = {{ bits }}u;
+// Codes per logical byte and the mask for one code. Physical uint8 storage uses
+// a u32 slot, so packing increases the number of codes returned by each load.
+const CODES_PER_BYTE: u32 = {{ codesPerByte }}u;
+const CODE_MASK: u32 = {{ codeMask }}u;
+{% for stream in (["q", "k", "v"] if not source.singleProjection else [source.singleProjection]) %}
+{{ matmul_nbits_packed_code(fn=stream ~ "_code", buffer=stream ~ "_b", kBlocks="KBLOCKS", blobSize="BLOB_SIZE", bits=bits) }}
+// Decode two consecutive reduction-axis codes from one stored word. An odd
+// offset would straddle words, so callers advance by two from an even start.
+fn {{ stream }}_code_pair(n: u32, block: u32, offset: u32) -> vec2<u32> {
+  let word = {{ stream }}_b[(n * KBLOCKS + block) * BLOB_SIZE + offset / CODES_PER_BYTE];
+  let shift = (offset % CODES_PER_BYTE) * BITS;
+  return vec2<u32>((word >> shift) & CODE_MASK, (word >> (shift + BITS)) & CODE_MASK);
+}
+{% endfor %}
+var<workgroup> reduction: array<f32, WG * ROW_TILE>;
+{% macro walk_block(codeFn, guarded) %}
+        for (var offset = lane * 2u; offset + 1u < BLOCK_SIZE; offset = offset + LANES * 2u) {
+          let k = k_base + offset;
+{% if guarded %}
+          if (k + 1u < K) {
+{% endif %}
+            let codes = {{ codeFn }}_pair(n, block, offset);
+            let code_lo = f32(codes.x) - ZERO;
+            let code_hi = f32(codes.y) - ZERO;
+{% for r in range(rowTile) %}
+            block_acc_{{ r }} = block_acc_{{ r }} + normed[base_{{ r }} + k] * code_lo
+                                                 + normed[base_{{ r }} + k + 1u] * code_hi;
+{% endfor %}
+{% if guarded %}
+          } else if (k < K) {
+            let code = f32({{ codeFn }}(n, block, offset)) - ZERO;
+{% for r in range(rowTile) %}
+            block_acc_{{ r }} = block_acc_{{ r }} + normed[base_{{ r }} + k] * code;
+{% endfor %}
+          }
+{% endif %}
+        }
+{%- endmacro %}
+{% macro project(codeFn, scalesBuffer) %}
+    for (var block = 0u; block < KBLOCKS; block = block + 1u) {
+      let scale = f32({{ scalesBuffer }}[n * KBLOCKS + block]);
+      let k_base = block * BLOCK_SIZE;
+{% for r in range(rowTile) %}
+      var block_acc_{{ r }} = 0.0;
+{% endfor %}
+      // Each trip handles two codes. BLOCK_SIZE is even for every admitted
+      // packing, so paired trips cover a full block. Only a final partial block
+      // needs bounds checks; the branch is workgroup-uniform.
+      if (k_base + BLOCK_SIZE <= K) {
+{{ walk_block(codeFn, false) }}
+      } else {
+{{ walk_block(codeFn, true) }}
+      }
+{% for r in range(rowTile) %}
+      acc_{{ r }} = acc_{{ r }} + block_acc_{{ r }} * scale;
+{% endfor %}
+    }
+{%- endmacro %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row0 = wg.y * ROW_TILE;
+  let tid = lid.x;
+  let column = tid / LANES;
+  let lane = tid % LANES;
+{% for r in range(rowTile) %}
+  // Rows past the end of the batch clamp onto the last real row; their
+  // accumulators are finite and the store guard drops them.
+  let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
+{% endfor %}
+{% if source.singleProjection %}
+  let n = wg.x * TILE_N + column;
+{% for r in range(rowTile) %}
+  var acc_{{ r }} = 0.0;
+{% endfor %}
+  if (n < N) {
+{{ project(source.singleProjection ~ "_code", source.singleProjection ~ "_scales") }}
+  }
+{% else %}
+  // Tile index selects the projection, so every thread in this workgroup takes
+  // the same arm and the reduction below stays uniform.
+  let tile = wg.x;
+  var projection = 2u;
+  var local_tile = tile - Q_TILES - KV_TILES;
+  if (tile < Q_TILES) {
+    projection = 0u;
+    local_tile = tile;
+  } else if (tile < Q_TILES + KV_TILES) {
+    projection = 1u;
+    local_tile = tile - Q_TILES;
+  }
+  let n = local_tile * TILE_N + column;
+  let limit = select(NKV, NQ, projection == 0u);
+{% for r in range(rowTile) %}
+  var acc_{{ r }} = 0.0;
+{% endfor %}
+  if (n < limit) {
+    if (projection == 0u) {
+{{ project("q_code", "q_scales") }}
+    } else if (projection == 1u) {
+{{ project("k_code", "k_scales") }}
+    } else {
+{{ project("v_code", "v_scales") }}
+    }
+  }
+{% endif %}
+{% for r in range(rowTile) %}
+  reduction[{{ r }}u * WG + tid] = acc_{{ r }};
+{% endfor %}
+  workgroupBarrier();
+  // Fold within each column's own LANES-wide group, one group per staged row.
+  // LANES is a power of two, so lane + stride never leaves the group.
+  for (var stride = LANES / 2u; stride > 0u; stride = stride / 2u) {
+    if (lane < stride) {
+{% for r in range(rowTile) %}
+      reduction[{{ r }}u * WG + tid] = reduction[{{ r }}u * WG + tid] + reduction[{{ r }}u * WG + tid + stride];
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% if source.singleProjection %}
+  if (lane == 0u && n < N) {
+{% for r in range(rowTile) %}
+    if (row0 + {{ r }}u < ROWS) {
+      {{ source.singleProjection }}[(row0 + {{ r }}u) * N + n] = {{ scalar }}(reduction[{{ r }}u * WG + tid]);
+    }
+{% endfor %}
+  }
+{% else %}
+  if (lane == 0u && n < limit) {
+{% for r in range(rowTile) %}
+    if (row0 + {{ r }}u < ROWS) {
+      let value_{{ r }} = {{ scalar }}(reduction[{{ r }}u * WG + tid]);
+      if (projection == 0u) {
+        q[(row0 + {{ r }}u) * NQ + n] = value_{{ r }};
+      } else if (projection == 1u) {
+        k[(row0 + {{ r }}u) * NKV + n] = value_{{ r }};
+      } else {
+        v[(row0 + {{ r }}u) * NKV + n] = value_{{ r }};
+      }
+    }
+{% endfor %}
+  }
+{% endif %}
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,663 @@

+{
+  "op": "com.microsoft.MatMulNBitsQkv",
+  "fixtureArrays": {
+    "norm_gqa_input_vBT": [90, 17, 203, 156, 64, 241, 112, 38, 175, 229, 83, 11, 198, 147, 52, 220, 105]
+  },
+  "cases": [
+    {
+      "name": "norm_gqa",
+      "provenance": {
+        "notes": "Grouped-query shape: Nq is twice Nkv, so the three projections have two different column counts."
+      },
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [3, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "norm_rows_past_one_tile",
+      "provenance": {
+        "notes": "Ten activation rows against an eight-row projection tile: two row groups where the second carries only two real rows, so the store guard has to drop the six that clamped onto the last row."
+      },
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [10, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [10, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [10, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [10, 4], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "skip_gqa",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32, "epsilon": 0.00001 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [3, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "skipsum_gqa",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [3, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "residualT": { "dtype": "float32", "shape": [3, 32], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "decode_single_row",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [1, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [1, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [1, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "residualT": { "dtype": "float32", "shape": [1, 32], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "equal_nq_nkv",
+      "provenance": {
+        "notes": "Multi-head shape where Nq equals Nkv; the tile grid still lays Q, K and V out in sequence."
+      },
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 8, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "tiles_not_aligned",
+      "provenance": {
+        "notes": "Neither column count is a multiple of the column tile, so both projections have a partial last tile."
+      },
+      "attrs": { "K": 32, "Nq": 5, "Nkv": 3, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [5, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [5, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [3, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [3, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [3, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [3, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [2, 5], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [2, 3], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [2, 3], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "tail_block_partial_k",
+      "provenance": { "notes": "K is not a multiple of block_size, so the final block is partly out of range." },
+      "attrs": { "K": 20, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 20],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [20],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "rank3_activation",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [1, 4, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [1, 4, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [1, 4, 4], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "two_quant_blocks",
+      "provenance": { "notes": "K=64 covers two of the provider's supported 32-element quantization blocks." },
+      "attrs": { "K": 64, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [2, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 2, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [8, 2],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 2, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 2, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [4, 2],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [2, 8], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "residualT": { "dtype": "float32", "shape": [2, 64], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "f16_skipsum",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [2, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [2, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float16",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float16",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float16",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float16",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float16", "shape": [2, 8], "tolerance": 0.002, "relTolerance": 0.01 },
+        "kT": { "dtype": "float16", "shape": [2, 4], "tolerance": 0.002, "relTolerance": 0.01 },
+        "vT": { "dtype": "float16", "shape": [2, 4], "tolerance": 0.002, "relTolerance": 0.01 },
+        "residualT": { "dtype": "float16", "shape": [2, 32], "tolerance": 0.002, "relTolerance": 0.002 }
+      }
+    },
+    {
+      "name": "f16_norm",
+      "attrs": { "K": 32, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float16",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 1, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float16",
+          "shape": [8, 1],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "kScalesT": {
+          "dtype": "float16",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 1, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float16",
+          "shape": [4, 1],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float16", "shape": [3, 8], "tolerance": 0.002, "relTolerance": 0.01 },
+        "kT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 },
+        "vT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 }
+      }
+    }
+  ]
+}