sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +29 -25
build/webgpu/bench.json +19 -15
build/webgpu/manifest.json +149 -358
build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja +5 -10
build/webgpu/metadata.json +18 -8
build/webgpu/qkv-projection.wgsl.jinja +213 -29
build/webgpu/test.json +166 -2

README.md CHANGED Viewed

@@ -18,26 +18,26 @@ See the [ONNX Runtime `MatMulNBitsQkv` contrib-operator spec](https://github.com
 ## 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
@@ -45,13 +45,13 @@ 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
@@ -62,7 +62,7 @@ Attributes and default values (overridable per request):
 ## 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
@@ -71,10 +71,14 @@ Attributes and default values (overridable per request):
 ## 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.
@@ -83,8 +87,8 @@ 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] },
@@ -93,7 +97,7 @@ const { qT, kT, vT } = await kernel({
   vScalesT: { data: vScalesTData, shape: [3, 1] },
 }, {
   attrs: {
-    K: 32,
     Nq: 5,
     Nkv: 3,
     block_size: 32,

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `aT` | `A` | `T1` | — | — | Shared activation of rank 2 `(M, K)` or rank 3 `(batch, sequence, K)`; only the last axis is the reduction axis. | required |
+| `skipT` | `skip` | `T1` | — | — | Residual added to A before the normalization, with A's shape. | optional |
+| `normScaleT` | `norm_scale` | `T1` | `1` | — | Simplified-layer-normalization (RMS) gain of shape `[K]`. | required |
+| `qBT` | `q_B` | `T2` | `3` | — | Bit-packed uint8 Q weights of shape `(Nq, k_blocks, blob_size)`. Bound in the packed storage layout: four blob bytes per u32 word. | required |
+| `qScalesT` | `q_scales` | `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 |
+| `kBT` | `k_B` | `T2` | `3` | — | Bit-packed K weights of shape `(Nkv, k_blocks, blob_size)`. Bound in the packed storage layout: four blob bytes per u32 word. | required |
+| `kScalesT` | `k_scales` | `T1` | `2` | — | Per-block K scales of shape `(Nkv, k_blocks)`. | required |
+| `vBT` | `v_B` | `T2` | `3` | — | Bit-packed V weights of shape `(Nkv, k_blocks, blob_size)`. Bound in the packed storage layout: four blob bytes per u32 word. | required |
+| `vScalesT` | `v_scales` | `T1` | `2` | — | Per-block V scales of shape `(Nkv, k_blocks)`. | required |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `qT` | `Q` | `T1` | same as `aT` | derived | Query projection: A's leading axes with a trailing Nq. | required |
+| `kT` | `K` | `T1` | same as `aT` | derived | Key projection: A's leading axes with a trailing Nkv. | required |
+| `vT` | `V` | `T1` | same as `aT` | derived | Value projection: A's leading axes with a trailing Nkv. | required |
+| `residualT` | `input_skip_bias_sum` | `T1` | same as `aT` | same as `aT` | The residual sum A + skip, with A's shape. Requires the skip input. | optional |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `K` | — | Input feature dimension shared by the normalized input and all projection weights. |
 | `Nkv` | — | Output feature dimension shared by the K and V projections. |
+| `Nq` | — | Output feature dimension of the Q projection. |
+| `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. |
 | `block_size` | — | Size of each quantization block along K; only 32 is supported. |
+| `epsilon` | `9.999999974752427e-7` | Epsilon used by the simplified layer-normalization reduction. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.
 const kernel = await getKernel("webgpu-kernels/com.microsoft.MatMulNBitsQkv", { version: 1 });
 const { qT, kT, vT } = await kernel({
+  aT: { data: aTData, shape: [1, 21] },
+  normScaleT: { data: normScaleTData, shape: [21] },
   qBT: { data: qBTData, shape: [5, 1, 16] },
   qScalesT: { data: qScalesTData, shape: [5, 1] },
   kBT: { data: kBTData, shape: [3, 1, 16] },
   vScalesT: { data: vScalesTData, shape: [3, 1] },
 }, {
   attrs: {
+    K: 21,
     Nq: 5,
     Nkv: 3,
     block_size: 32,

build/webgpu/bench.json CHANGED Viewed

@@ -1,20 +1,20 @@
 {
-  "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": {
@@ -22,13 +22,14 @@
         "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 },
@@ -75,13 +76,14 @@
         "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 },
@@ -128,13 +130,14 @@
         "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 },
@@ -181,13 +184,14 @@
         "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 },
@@ -234,13 +238,13 @@
         "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 },
@@ -287,7 +291,7 @@
         "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" }] }
     }
   ]
 }

 {
   "tunableSpace": { "TILE_N": [4, 8, 16], "LANES": [4, 8, 16] },
   "cases": [
     {
       "name": "qkv-q4-decode-k2048-nq2048-nkv512",
+      "tunableSpace": {},
       "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, "min": 0, "max": 256 },
         "qScalesT": { "shape": [2048, 64], "dtype": "float32", "dist": "normal", "seed": 9104, "scale": 0.05 },
+        "kBT": { "shape": [512, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 9105, "min": 0, "max": 256 },
         "kScalesT": { "shape": [512, 64], "dtype": "float32", "dist": "normal", "seed": 9106, "scale": 0.05 },
+        "vBT": { "shape": [512, 64, 16], "dtype": "uint8", "dist": "uniform", "seed": 9107, "min": 0, "max": 256 },
         "vScalesT": { "shape": [512, 64], "dtype": "float32", "dist": "normal", "seed": 9108, "scale": 0.05 }
       },
       "outputs": {
         "kT": { "shape": [1, 512], "dtype": "float32" },
         "vT": { "shape": [1, 512], "dtype": "float32" }
       },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "(2048 + 512 + 512) * 64 * 16" }] }
     },
     {
       "name": "qkv-q4-llama-decode-k4096-nq4096-nkv4096",
+      "tunableSpace": {},
       "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 the packed blob: the weight ports bind in the packed storage layout (four blob bytes per u32 word), so a 4-bit code costs half a byte of traffic and the projection moves its on-disk size."
       },
       "vars": { "dtype": "float32", "tokens": 1, "K": 4096, "Nq": 4096, "Nkv": 4096 },
       "attrs": { "K": 4096, "Nq": 4096, "Nkv": 4096, "block_size": 32 },
         "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",
+      "tunableSpace": {},
       "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 the packed blob: the weight ports bind in the packed storage layout (four blob bytes per u32 word), so a 4-bit code costs half a byte of traffic and the projection moves its on-disk size."
       },
       "vars": { "dtype": "float32", "tokens": 1, "K": 4096, "Nq": 4096, "Nkv": 1024 },
       "attrs": { "K": 4096, "Nq": 4096, "Nkv": 1024, "block_size": 32 },
         "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",
+      "tunableSpace": {},
       "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 the packed blob: the weight ports bind in the packed storage layout (four blob bytes per u32 word), so a 4-bit code costs half a byte of traffic and the projection moves its on-disk size."
       },
       "vars": { "dtype": "float32", "tokens": 1, "K": 2304, "Nq": 2048, "Nkv": 1024 },
       "attrs": { "K": 2304, "Nq": 2048, "Nkv": 1024, "block_size": 32 },
         "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",
+      "tunableSpace": {},
       "preset": "model",
       "provenance": {
+        "notes": "Phi-3 class defaults (hidden_size 3072, num_attention_heads 32, no GQA). Bytes are the packed blob: the weight ports bind in the packed storage layout (four blob bytes per u32 word), so a 4-bit code costs half a byte of traffic and the projection moves its on-disk size."
       },
       "vars": { "dtype": "float32", "tokens": 1, "K": 3072, "Nq": 3072, "Nkv": 3072 },
       "attrs": { "K": 3072, "Nq": 3072, "Nkv": 3072, "block_size": 32 },
         "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 the packed blob: the weight ports bind in the packed storage layout (four blob bytes per u32 word), so a 4-bit code costs half a byte of traffic and the projection moves its on-disk size."
       },
       "vars": { "dtype": "float32", "tokens": 8, "K": 4096, "Nq": 4096, "Nkv": 4096 },
       "attrs": { "K": 4096, "Nq": 4096, "Nkv": 4096, "block_size": 32 },
         "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 CHANGED Viewed

@@ -2,92 +2,37 @@
   "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 },
@@ -98,22 +43,13 @@
     "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)",
@@ -131,22 +67,21 @@
     "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",
@@ -156,378 +91,234 @@
     "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" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "MatMulNBitsQkv",
   "sinceVersion": 1,
+  "inputs": {
+    "aT": { "onnx": "A", "dtype": "T1" },
+    "skipT": { "onnx": "skip", "dtype": "T1", "optional": true },
+    "normScaleT": { "onnx": "norm_scale", "dtype": "T1", "rank": 1 },
+    "qBT": { "onnx": "q_B", "dtype": "T2", "rank": 3, "layout": "packed" },
+    "qScalesT": { "onnx": "q_scales", "dtype": "T1", "rank": 2 },
+    "kBT": { "onnx": "k_B", "dtype": "T2", "rank": 3, "layout": "packed" },
+    "kScalesT": { "onnx": "k_scales", "dtype": "T1", "rank": 2 },
+    "vBT": { "onnx": "v_B", "dtype": "T2", "rank": 3, "layout": "packed" },
+    "vScalesT": { "onnx": "v_scales", "dtype": "T1", "rank": 2 }
+  },
+  "outputs": {
+    "qT": { "onnx": "Q", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nq]" },
+    "kT": { "onnx": "K", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nkv]" },
+    "vT": { "onnx": "V", "dtype": "T1", "rank": "ranks.aT", "shape": "shapes.aT[:-1] + [attrs.Nkv]" },
+    "residualT": {
+      "onnx": "input_skip_bias_sum",
       "dtype": "T1",
       "rank": "ranks.aT",
       "optional": true,
+      "shape": "shapes.aT"
     }
+  },
+  "attributes": {
+    "accuracy_level": { "default": 0 },
+    "bits": { "default": 4 },
+    "epsilon": { "default": 9.999999974752427e-7 },
+    "K": {},
+    "Nq": {},
+    "Nkv": {},
+    "block_size": {}
   },
   "attributeConstraints": {
     "K": { "required": true },
     "block_size": { "required": true, "values": [32] }
   },
   "typeConstraints": { "T1": ["float32", "float16"], "T2": ["uint8"] },
+  "tunables": {
+    "TILE_N": { "default": 8 },
+    "LANES": { "default": 8 },
+    "NORM_WORKGROUP_SIZE": { "default": 128 },
+    "ROW_TILE": { "default": 8 },
+    "DECODE_WORKGROUP_SIZE": { "default": 64 }
   },
   "derive": {
     "aRows": "numel(shapes.aT) / max(1, attrs.K)",
     "rowTile": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)",
     "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",
+    "decodeWalk": "aRows <= 1",
+    "decodeCols": "4",
+    "gemvWalk": "rowTile == 1 and blobSize % 16 == 0",
+    "decodeActVec4": "gemvWalk and attrs.K % attrs.block_size == 0",
+    "tileCols": "decodeCols if decodeWalk else tunables.TILE_N",
+    "decodeWorkgroupOk": "tunables.DECODE_WORKGROUP_SIZE >= 4 and pow2ceil(tunables.DECODE_WORKGROUP_SIZE) == tunables.DECODE_WORKGROUP_SIZE and tunables.DECODE_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.DECODE_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX",
+    "projectionTiles": "ceilDiv(attrs.Nq, tileCols) + 2 * ceilDiv(attrs.Nkv, tileCols)",
+    "dispatchFits": "decodeWorkgroupOk and projectionTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and aRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) 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",
     "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
     "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
     "K": "attrs.K",
     "nq": "attrs.Nq",
     "nkv": "attrs.Nkv",
     "blockSize": "attrs.block_size",
     "bits": "attrs.bits",
     "defaultZero": "\"8.0\"",
     "tileN": "tunables.TILE_N",
     "lanes": "tunables.LANES",
     "hasSkip": "present.skipT",
     "writeResidual": "present.residualT",
     "K_LEN": "attrs.K",
+    "rowCount": "aRows",
+    "decodeNCols": "decodeCols",
+    "actVec4": "decodeActVec4",
+    "weightElement": "\"vec4<u32>\" if gemvWalk else \"u32\"",
+    "normedElement": "\"vec4<f32>\" if decodeActVec4 else \"f32\"",
+    "decodeWorkgroupSize": "tunables.DECODE_WORKGROUP_SIZE",
+    "useSubgroups": "device.features.has(\"subgroups\")"
   },
+  "when": ["dispatchFits", "qkvShapeOk"],
+  "bindings": {
+    "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "norm_scale": { "arg": "normScaleT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$K_LEN" },
+    "normed": { "scratch": "normedA", "buffer": "storage", "elementType": "f32" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "rows", "type": "u32", "value": "aRows" }] },
+    "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "residual": { "arg": "residualT", "buffer": "storage", "elementType": "$aScalar" },
+    "normed_2": {
+      "scratch": "normedA",
+      "name": "normed",
+      "buffer": "read-only-storage",
+      "elementType": "$normedElement"
+    },
+    "q_b": { "arg": "qBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
+    "q_scales": { "arg": "qScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "k_b": { "arg": "kBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
+    "k_scales": { "arg": "kScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "v_b": { "arg": "vBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
+    "v_scales": { "arg": "vScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "q": { "arg": "qT", "buffer": "storage", "elementType": "$aScalar" },
+    "k": { "arg": "kT", "buffer": "storage", "elementType": "$aScalar" },
+    "v": { "arg": "vT", "buffer": "storage", "elementType": "$aScalar" }
   },
   "variants": [
     {
       "id": "norm",
       "priority": 20,
+      "when": ["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": ["a", "norm_scale", "normed", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "main",
           "name": "MatMulNBitsQkv.Projection",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"],
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "split_norm",
       "priority": 10,
+      "when": ["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": ["a", "norm_scale", "normed", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "q",
           "name": "MatMulNBitsQkv.ProjectionQ",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"q\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "q"],
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "k",
           "name": "MatMulNBitsQkv.ProjectionK",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"k\"" },
+          "bindings": ["normed_2", "k_b", "k_scales", "k"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "v",
           "name": "MatMulNBitsQkv.ProjectionV",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"v\"" },
+          "bindings": ["normed_2", "v_b", "v_scales", "v"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "skip",
       "priority": 20,
+      "when": ["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": ["a", "skip", "norm_scale", "normed", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "main",
           "name": "MatMulNBitsQkv.Projection",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"],
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "split_skip",
       "priority": 10,
+      "when": ["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": ["a", "skip", "norm_scale", "normed", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "q",
           "name": "MatMulNBitsQkv.ProjectionQ",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"q\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "q"],
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "k",
           "name": "MatMulNBitsQkv.ProjectionK",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"k\"" },
+          "bindings": ["normed_2", "k_b", "k_scales", "k"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "v",
           "name": "MatMulNBitsQkv.ProjectionV",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"v\"" },
+          "bindings": ["normed_2", "v_b", "v_scales", "v"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "skipsum",
       "priority": 20,
+      "when": ["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": ["a", "skip", "norm_scale", "normed", "residual", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "main",
           "name": "MatMulNBitsQkv.Projection",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "k_b", "k_scales", "v_b", "v_scales", "q", "k", "v"],
+          "dispatch": { "x": "projectionTiles", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "split_skipsum",
       "priority": 10,
+      "when": ["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": ["a", "skip", "norm_scale", "normed", "residual", "params"],
+          "dispatch": { "x": "min(aRows, 65535)", "y": "ceilDiv(aRows, 65535)", "z": 1 }
         },
         {
           "id": "q",
           "name": "MatMulNBitsQkv.ProjectionQ",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"q\"" },
+          "bindings": ["normed_2", "q_b", "q_scales", "q"],
+          "dispatch": { "x": "ceilDiv(attrs.Nq, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "k",
           "name": "MatMulNBitsQkv.ProjectionK",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"k\"" },
+          "bindings": ["normed_2", "k_b", "k_scales", "k"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         },
         {
           "id": "v",
           "name": "MatMulNBitsQkv.ProjectionV",
+          "shader": "qkv-projection.wgsl.jinja",
+          "derive": { "singleProjection": "\"v\"" },
+          "bindings": ["normed_2", "v_b", "v_scales", "v"],
+          "dispatch": { "x": "ceilDiv(attrs.Nkv, tileCols)", "y": "rowGroups" },
+          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

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

@@ -1,14 +1,10 @@
-{% 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 }};
@@ -88,11 +84,10 @@ fn row_value(index: u32) -> f32 {
 @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;
   }

 {{ env.wgsl.resourceDeclarations }}
+// Fused RMS-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. The following projection therefore consumes the
+// normalized values without an intervening storage-type narrowing.
 const HIDDEN: u32 = {{ hidden }}u;
 const WG: u32 = {{ workgroupSize }}u;
 const EPSILON: f32 = {{ epsilon }};
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits past the
+  // per-axis dispatch fold width. Reduces to wg.x when the dispatch does not fold.
+  let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (row >= params.rows) {
     return;
   }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,29 @@
 {
   "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" }
 }

 {
   "name": "com.microsoft.MatMulNBitsQkv",
+  "id": "_com_microsoft_matmulnbitsqkv_webgpu_f6dfd6b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "58HIjRCefPnXYJoowt4ER40ZwXZBaytxqPwmLff2bUM=",
+      "manifest.json": "fAGWh8u4euc77NPIK/A78FwhhxsO0TXSThumgnholyE=",
+      "matmul-nbits-fused-rms-norm.wgsl.jinja": "4lOdB+RprQh3iv29i6RV8UWkn8S5y1aK6Te5nJpxuEk=",
+      "qkv-projection.wgsl.jinja": "M0SyxodhZFrp/YOpcBTegsqP0M2Vz98ZwSMxSDDa8ew=",
+      "test.json": "tpaBr6bhveKXElA96j2p78/ivs6jAtlZyTmSxQbbb5w="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "norm": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"],
+      "split_norm": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"],
+      "skip": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"],
+      "split_skip": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"],
+      "skipsum": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"],
+      "split_skipsum": ["matmul-nbits-fused-rms-norm.wgsl.jinja", "qkv-projection.wgsl.jinja"]
+    }
+  }
 }

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

@@ -1,30 +1,28 @@
 {% 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
@@ -42,45 +40,231 @@ enable f16;
 // 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) %}
@@ -140,14 +324,14 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   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
@@ -194,11 +378,11 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     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 %}
   }
@@ -218,4 +402,4 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% endfor %}
   }
 {% endif %}
-}

 {% 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 = (n * {{ kBlocks }} + block) * {{ blobSize }} + (offset >> 2u);
+  let shift = (byte_index & 3u) * 8u + (offset & 3u) * 2u;
+  return ({{ buffer }}[byte_index >> 2u] >> shift) & 3u;
 {% elif bits == 4 %}
+  let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + (offset >> 1u);
+  let shift = (byte_index & 3u) * 8u + (offset & 1u) * 4u;
+  return ({{ buffer }}[byte_index >> 2u] >> shift) & 15u;
 {% else %}
+  let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
+  return ({{ buffer }}[byte_index >> 2u] >> ((byte_index & 3u) * 8u)) & 255u;
 {% endif %}
 }
 {%- endmacro %}
+{% if gemvWalk and useSubgroups %}
+enable subgroups;
 {% 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 singleProjection %}
 // This specialization binds and computes one projection.
 {% else %}
 // All three projections read the same normalized row, so one dispatch covers
 // A workgroup also covers ROW_TILE activation rows, reusing each unpacked code
 // across their accumulators.
 const K: u32 = {{ K }}u;
+{% if singleProjection == "q" %}
 const N: u32 = {{ nq }}u;
+{% elif singleProjection %}
 const N: u32 = {{ nkv }}u;
 {% else %}
 const NQ: u32 = {{ nq }}u;
 const NKV: u32 = {{ nkv }}u;
 {% endif %}
+{% if not gemvWalk %}
 const BLOCK_SIZE: u32 = {{ blockSize }}u;
+{% endif %}
 const KBLOCKS: u32 = {{ kBlocks }}u;
 const BLOB_SIZE: u32 = {{ blobSize }}u;
+{% if gemvWalk %}
+const N_COLS: u32 = {{ decodeNCols }}u;
+const WG: u32 = {{ decodeWorkgroupSize }}u;
+const ROWS: u32 = {{ rowCount }}u;
+const ZERO: f32 = {{ defaultZero }};
+{% if not singleProjection %}
+const Q_TILES: u32 = (NQ + N_COLS - 1u) / N_COLS;
+const KV_TILES: u32 = (NKV + N_COLS - 1u) / N_COLS;
+{% endif %}
+{% else %}
 const TILE_N: u32 = {{ tileN }}u;
 const LANES: u32 = {{ lanes }}u;
 const ROW_TILE: u32 = {{ rowTile }}u;
+const ROWS: u32 = {{ rowCount }}u;
 const WG: u32 = TILE_N * LANES;
 const ZERO: f32 = {{ defaultZero }};
+{% if not singleProjection %}
 const Q_TILES: u32 = (NQ + TILE_N - 1u) / TILE_N;
 const KV_TILES: u32 = (NKV + TILE_N - 1u) / TILE_N;
 {% endif %}
+{% endif %}
+{% if not gemvWalk %}
 const BITS: u32 = {{ bits }}u;
+{% endif %}
+// Codes per blob byte and the mask for one code. The blob is bound in the
+// packed storage layout -- four blob bytes per u32 word -- so every load
+// returns 4 * CODES_PER_BYTE codes.
 const CODES_PER_BYTE: u32 = {{ codesPerByte }}u;
 const CODE_MASK: u32 = {{ codeMask }}u;
+{% if gemvWalk %}
+// One vec4<u32> of the packed blob is sixteen stored bytes, so it carries
+// 16 * CODES_PER_BYTE codes -- and a whole quantization block is exactly
+// BLOB_SIZE / 16 of them (this walk requires blobSize % 16 == 0). A vector
+// therefore never straddles two blocks, so its codes all share one scale, and
+// the flat vector index needs no division at all:
+//
+//   (n * KBLOCKS + block) * (BLOB_SIZE / 16) + slot  ==  n * VEC_GROUPS + group
+const VEC_PER_BLOCK: u32 = BLOB_SIZE / 16u;
+const VEC_GROUPS: u32 = KBLOCKS * VEC_PER_BLOCK;
+const CODES_PER_VEC: u32 = 16u * CODES_PER_BYTE;
+{% endif %}
+{% for stream in (["q", "k", "v"] if not singleProjection else [singleProjection]) %}
+{% if not gemvWalk %}
 {{ 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 byte. An odd
+// offset would straddle bytes, so callers advance by two from an even start.
 fn {{ stream }}_code_pair(n: u32, block: u32, offset: u32) -> vec2<u32> {
+  let byte_index = (n * KBLOCKS + block) * BLOB_SIZE + offset / CODES_PER_BYTE;
+  let shift = (byte_index & 3u) * 8u + (offset % CODES_PER_BYTE) * BITS;
+  let word = {{ stream }}_b[byte_index >> 2u];
   return vec2<u32>((word >> shift) & CODE_MASK, (word >> (shift + BITS)) & CODE_MASK);
 }
+{% endif %}
 {% endfor %}
+{% if gemvWalk %}
+var<workgroup> partials: array<vec4<f32>, WG>;
+{% set codesPerWord = 4 * codesPerByte %}
+{% macro gemv_project(weights, scalesBuffer) %}
+    for (var g = tid; g < VEC_GROUPS; g = g + WG) {
+      let block = g / VEC_PER_BLOCK;
+{% if actVec4 %}
+      // K is a whole number of blocks here, so the group lands entirely inside the
+      // row and the activations come from CODES_PER_VEC / 4 aligned vector loads.
+      let vbase = base4 + g * (CODES_PER_VEC / 4u);
+{% for v in range(codesPerWord) %}
+      let av{{ v }} = normed[vbase + {{ v }}u];
+{% endfor %}
+{% for v in range(codesPerWord) %}
+{% for c in range(4) %}
+      let a{{ v * 4 + c }} = av{{ v }}.{{ ["x", "y", "z", "w"][c] }};
+{% endfor %}
+{% endfor %}
+{% else %}
+      let k0 = g * CODES_PER_VEC;
+{% for j in range(16 * codesPerByte) %}
+      // A trailing partial block reaches past K; those codes contribute zero.
+      let a{{ j }} = select(0.0, normed[base + min(k0 + {{ j }}u, K - 1u)], k0 + {{ j }}u < K);
+{% endfor %}
+{% endif %}
+      let asum = {% for j in range(16 * codesPerByte) %}{{ " + " if j > 0 else "" }}a{{ j }}{% endfor %};
+{% for c in range(4) %}
+{% set comp = ["x", "y", "z", "w"][c] %}
+{% if c == 0 %}
+      {
+{% else %}
+      if (col_base + {{ c }}u < limit) {
+{% endif %}
+        let n = col_base + {{ c }}u;
+        let words = {{ weights }}[n * VEC_GROUPS + g];
+        let scale = f32({{ scalesBuffer }}[n * KBLOCKS + block]);
+        var dot = 0.0;
+{% for w in range(4) %}
+{% for h in range(codesPerWord) %}
+        dot = dot + a{{ w * codesPerWord + h }} * f32((words.{{ ["x", "y", "z", "w"][w] }} >> {{ h * bits }}u) & CODE_MASK);
+{% endfor %}
+{% endfor %}
+        acc.{{ comp }} = acc.{{ comp }} + (dot - ZERO * asum) * scale;
+      }
+{% endfor %}
+    }
+{%- endmacro %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
+        @builtin(subgroup_invocation_id) sgLane: u32, @builtin(subgroup_size) sgSize: u32{% endif %}) {
+  let tid = lid.x;
+  // Rows past the end of the batch clamp onto the last real row; the store
+  // guard drops them.
+  let row = min(wg.y, ROWS - 1u);
+  let base = row * K;
+{% if actVec4 %}
+  let base4 = base / 4u;
+{% endif %}
+{% if singleProjection %}
+  let col_base = wg.x * N_COLS;
+  let limit = N;
+{% 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 col_base = local_tile * N_COLS;
+  let limit = select(NKV, NQ, projection == 0u);
+{% endif %}
+  // Whole workgroups past the last column return before any barrier.
+  if (col_base >= limit) {
+    return;
+  }
+  var acc = vec4<f32>(0.0);
+{% if singleProjection %}
+{{ gemv_project(singleProjection ~ "_b", singleProjection ~ "_scales") }}
+{% else %}
+  if (projection == 0u) {
+{{ gemv_project("q_b", "q_scales") }}
+  } else if (projection == 1u) {
+{{ gemv_project("k_b", "k_scales") }}
+  } else {
+{{ gemv_project("v_b", "v_scales") }}
+  }
+{% endif %}
+{% if useSubgroups %}
+  // Subgroup fold: one collective, then the WG / subgroup-size per-subgroup
+  // partials fold once through workgroup memory. The lanes of a subgroup are
+  // contiguous in local_invocation_id; the projection branches above are
+  // workgroup-uniform, so the collective runs in uniform control flow.
+  let sgSum = subgroupAdd(acc);
+  if (sgLane == 0u) {
+    partials[tid / sgSize] = sgSum;
+  }
+  workgroupBarrier();
+  if (tid == 0u && wg.y < ROWS) {
+    let subgroupCount = WG / sgSize;
+    var total = partials[0];
+    for (var i = 1u; i < subgroupCount; i = i + 1u) {
+      total = total + partials[i];
+    }
+{% else %}
+  partials[tid] = acc;
+  workgroupBarrier();
+  var stride = WG / 2u;
+  loop {
+    if (stride == 0u) {
+      break;
+    }
+    if (tid < stride) {
+      partials[tid] = partials[tid] + partials[tid + stride];
+    }
+    stride = stride / 2u;
+    workgroupBarrier();
+  }
+  if (tid == 0u && wg.y < ROWS) {
+    let total = partials[0];
+{% endif %}
+{% for c in range(4) %}
+{% set comp = ["x", "y", "z", "w"][c] %}
+{% if c == 0 %}
+    {
+{% else %}
+    if (col_base + {{ c }}u < limit) {
+{% endif %}
+      let value = {{ scalar }}(total.{{ comp }});
+{% if singleProjection %}
+      {{ singleProjection }}[row * N + col_base + {{ c }}u] = value;
+{% else %}
+      if (projection == 0u) {
+        q[row * NQ + col_base + {{ c }}u] = value;
+      } else if (projection == 1u) {
+        k[row * NKV + col_base + {{ c }}u] = value;
+      } else {
+        v[row * NKV + col_base + {{ c }}u] = value;
+      }
+{% endif %}
+    }
+{% endfor %}
+  }
+}
+{% else %}
 var<workgroup> reduction: array<f32, WG * ROW_TILE>;
 {% macro walk_block(codeFn, guarded) %}
   let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
 {% endfor %}
+{% if singleProjection %}
   let n = wg.x * TILE_N + column;
 {% for r in range(rowTile) %}
   var acc_{{ r }} = 0.0;
 {% endfor %}
   if (n < N) {
+{{ project(singleProjection ~ "_code", singleProjection ~ "_scales") }}
   }
 {% else %}
   // Tile index selects the projection, so every thread in this workgroup takes
     workgroupBarrier();
   }
+{% if singleProjection %}
   if (lane == 0u && n < N) {
 {% for r in range(rowTile) %}
     if (row0 + {{ r }}u < ROWS) {
+      {{ singleProjection }}[(row0 + {{ r }}u) * N + n] = {{ scalar }}(reduction[{{ r }}u * WG + tid]);
     }
 {% endfor %}
   }
 {% endfor %}
   }
 {% endif %}
+}{% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "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]
   },
@@ -495,7 +494,7 @@
     },
     {
       "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": {
@@ -658,6 +657,171 @@
         "kT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 },
         "vT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "norm_gqa_input_vBT": [90, 17, 203, 156, 64, 241, 112, 38, 175, 229, 83, 11, 198, 147, 52, 220, 105]
   },
     },
     {
       "name": "two_quant_blocks",
+      "provenance": { "notes": "K=64 spans two 32-element quantization blocks." },
       "attrs": { "K": 64, "Nq": 8, "Nkv": 4, "block_size": 32 },
       "inputs": {
         "aT": {
         "kT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 },
         "vT": { "dtype": "float16", "shape": [3, 4], "tolerance": 0.002, "relTolerance": 0.01 }
       }
+    },
+    {
+      "name": "decode_multi_trip_tails",
+      "provenance": {
+        "notes": "Single row at K past two 128-code trips of the reduction-partitioned walk, with a partial final block and column counts not aligned to the decode column group."
+      },
+      "attrs": { "K": 300, "Nq": 5, "Nkv": 3, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 300],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [300],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [5, 10, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float32",
+          "shape": [5, 10],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [3, 10, 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, 10],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [3, 10, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float32",
+          "shape": [3, 10],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "decode_odd_k",
+      "provenance": { "notes": "Single row with an odd K, so the last reduction pair has one live code." },
+      "attrs": { "K": 21, "Nq": 5, "Nkv": 3, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 21],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [21],
+          "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": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "kT": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "vT": { "dtype": "float32", "shape": [1, 3], "tolerance": 0.0001, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "decode_f16_two_trips",
+      "provenance": {
+        "notes": "Single float16 row over two trips of the reduction-partitioned walk with the skip input."
+      },
+      "attrs": { "K": 160, "Nq": 8, "Nkv": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [1, 160],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "normScaleT": {
+          "dtype": "float16",
+          "shape": [160],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "qBT": {
+          "dtype": "uint8",
+          "shape": [8, 5, 16],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "qScalesT": {
+          "dtype": "float16",
+          "shape": [8, 5],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "kBT": {
+          "dtype": "uint8",
+          "shape": [4, 5, 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, 5],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        },
+        "vBT": {
+          "dtype": "uint8",
+          "shape": [4, 5, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/norm_gqa_input_vBT" } }
+        },
+        "vScalesT": {
+          "dtype": "float16",
+          "shape": [4, 5],
+          "data": { "kind": "cycle", "values": [0.052, 0.038, 0.061, 0.029, 0.073] }
+        },
+        "skipT": {
+          "dtype": "float16",
+          "shape": [1, 160],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "qT": { "dtype": "float16", "shape": [1, 8], "tolerance": 0.02, "relTolerance": 0.02 },
+        "kT": { "dtype": "float16", "shape": [1, 4], "tolerance": 0.02, "relTolerance": 0.02 },
+        "vT": { "dtype": "float16", "shape": [1, 4], "tolerance": 0.02, "relTolerance": 0.02 }
+      }
     }
   ]
 }