sync 91d990483a17

Browse files

Files changed (14) hide show

README.md +31 -13
build/webgpu/bench.json +54 -13
build/webgpu/manifest.json +0 -0
build/webgpu/matmul-nbits-dp4a-quantize.wgsl.jinja +1 -2
build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja +93 -47
build/webgpu/matmul-nbits-q4-dp4a-prefill.wgsl.jinja +7 -4
build/webgpu/matmul-nbits-q4-prefill-tile4x4.wgsl.jinja +105 -93
build/webgpu/matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja +43 -22
build/webgpu/matmul-nbits-q4-prefill-tiled.wgsl.jinja +11 -16
build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja +29 -25
build/webgpu/matmul-nbits.wgsl.jinja +11 -16
build/webgpu/metadata.json +52 -14
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +122 -0
build/webgpu/test.json +306 -40

README.md CHANGED Viewed

@@ -18,19 +18,19 @@ See the [ONNX Runtime `MatMulNBits` contrib-operator spec](https://github.com/mi
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `A` | `aT` | `T1` | — | — | Float input matrix, not quantized. Rank 2 has shape `(M, K)` and rank 3 has shape `(batch, sequence, K)`; only the last axis is the reduction axis and the leading axes fold into the row count, so the ordinary activation needs no surrounding Reshape. | required |
-| `B` | `bT` | `uint8` | `3` | — | Bit-packed uint8 weight matrix of shape `(N, k_blocks, blob_size)`, where `k_blocks = ceil(K / block_size)` and `blob_size = block_size * bits / 8`. Codes are packed low-first along K. | required |
-| `scales` | `scalesT` | `T1` | `2` | — | Per-block dequantization scale factors of shape `(N, k_blocks)`, with the same dtype as `A`. | required |
-| `zero_points` | `zeroPointsT` | `T3` | `2` | — | Standard unpacked per-block zero points with shape `(N, k_blocks)` and the same dtype as `A`. Omission uses `2^(bits - 1)`. | optional |
-| `bias` | `biasT` | `T1` | `1` | — | Optional bias vector of shape `[N]` added to the output. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `Y` | `yT` | `T1` | same as `A` | derived; see description | Result of A multiplied by the dequantized weight matrix, with optional bias, same dtype and rank as A: the leading axes of A with a trailing N. | required |
 ## Attributes
@@ -38,10 +38,10 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
-| `accuracy_level` | `0` | Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8). |
-| `bits` | `4` | Bit width used to quantize B; this package supports 2, 4, and 8. |
 | `K` | — | Input feature dimension of the weight matrix. |
 | `N` | — | Output feature dimension of the weight matrix. |
 | `block_size` | — | Power-of-two quantization block size along K; it must be at least 16. |
 ## Type constraints
@@ -51,13 +51,26 @@ Attributes and default values (overridable per request):
 | `T1` | `float32`, `float16` |
 | `T3` | `float32`, `float16` |
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, 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
@@ -69,13 +82,18 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`matmul-nbits-q4-prefill-tiled.wgsl.jinja`](build/webgpu/matmul-nbits-q4-prefill-tiled.wgsl.jinja)
 - [`matmul-nbits-q4-sgmat.wgsl.jinja`](build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja)
 - [`matmul-nbits.wgsl.jinja`](build/webgpu/matmul-nbits.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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `aT` | `A` | `T1` | — | — | Float input matrix, not quantized. Rank 2 has shape `(M, K)` and rank 3 has shape `(batch, sequence, K)`; only the last axis is the reduction axis and the leading axes fold into the row count, so the ordinary activation needs no surrounding Reshape. | required |
+| `bT` | `B` | `uint8` | `3` | — | Bit-packed uint8 weight matrix of shape `(N, k_blocks, blob_size)`, where `k_blocks = ceil(K / block_size)` and `blob_size = block_size * bits / 8`. Codes are packed low-first along K. Bound in the packed storage layout: four blob bytes per u32 word, so the kernels stream the blob's own bytes rather than one widened word per byte. | required |
+| `scalesT` | `scales` | `T1` | `2` | — | Per-block dequantization scale factors of shape `(N, k_blocks)`, with the same dtype as `A`. | required |
+| `zeroPointsT` | `zero_points` | `T3` | `2` | — | Standard unpacked per-block zero points with shape `(N, k_blocks)` and the same dtype as `A`. Omission uses `2^(bits - 1)`. | optional |
+| `biasT` | `bias` | `T1` | `1` | — | Optional bias vector of shape `[N]` added to the output. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `yT` | `Y` | `T1` | same as `aT` | derived | Result of A multiplied by the dequantized weight matrix, with optional bias, same dtype and rank as A: the leading axes of A with a trailing N. | required |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `K` | — | Input feature dimension of the weight matrix. |
 | `N` | — | Output feature dimension of the weight matrix. |
+| `accuracy_level` | `0` | Minimum internal accuracy level: 0 (unset), 1 (float32), 2 (float16), 3 (bfloat16), or 4 (int8). |
+| `bits` | `4` | Bit width used to quantize B; this package supports 2, 4, and 8. |
 | `block_size` | — | Power-of-two quantization block size along K; it must be at least 16. |
 ## Type constraints
 | `T1` | `float32`, `float16` |
 | `T3` | `float32`, `float16` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `prefill_tiled_reg_vec4_splitk_default_zero` — Cuts the K reduction of the four-wide register-blocked prefill tile into power-of-two slices across dispatch.z, retaining at least 512 K values per slice, then sums the f32 partials and adds bias in a second pass.
+- `prefill_tiled_reg_vec4_default_zero` — The register-blocked prefill tile with four-wide activation loads and a 128-row tile above 256 rows: the activation slice of every K tile is staged with one vector load per lane instead of four scalar loads, and the taller tile halves the dequantization work per multiply-add.
+- `prefill_tiled_reg_vec4_splitk_zero_bias` — Cuts the K reduction of the four-wide register-blocked prefill tile into power-of-two slices across dispatch.z, retaining at least 512 K values per slice, then sums the f32 partials and adds bias in a second pass.
+- `prefill_tiled_reg_vec4_zero_bias` — The register-blocked prefill tile with four-wide activation loads and a 128-row tile above 256 rows: the activation slice of every K tile is staged with one vector load per lane instead of four scalar loads, and the taller tile halves the dequantization work per multiply-add.
+- `prefill_tiled_reg_vec4_splitk_zero_only` — Cuts the K reduction of the four-wide register-blocked prefill tile into power-of-two slices across dispatch.z, retaining at least 512 K values per slice, then sums the f32 partials and adds bias in a second pass.
+- `prefill_tiled_reg_vec4_zero_only` — The register-blocked prefill tile with four-wide activation loads and a 128-row tile above 256 rows: the activation slice of every K tile is staged with one vector load per lane instead of four scalar loads, and the taller tile halves the dequantization work per multiply-add.
+- `prefill_tiled_reg_vec4_splitk_bias_only` — Cuts the K reduction of the four-wide register-blocked prefill tile into power-of-two slices across dispatch.z, retaining at least 512 K values per slice, then sums the f32 partials and adds bias in a second pass.
+- `prefill_tiled_reg_vec4_bias_only` — The register-blocked prefill tile with four-wide activation loads and a 128-row tile above 256 rows: the activation slice of every K tile is staged with one vector load per lane instead of four scalar loads, and the taller tile halves the dequantization work per multiply-add.
 ## Device requirements
 Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 - [`matmul-nbits-q4-prefill-tiled.wgsl.jinja`](build/webgpu/matmul-nbits-q4-prefill-tiled.wgsl.jinja)
 - [`matmul-nbits-q4-sgmat.wgsl.jinja`](build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja)
 - [`matmul-nbits.wgsl.jinja`](build/webgpu/matmul-nbits.wgsl.jinja)
+- [`reduce-axis0-splitk-combine.wgsl.jinja`](build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja)
 ## 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.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.MatMulNBits",
   "tunableSpace": { "workgroupSize": [64, 128, 256] },
   "cases": [
     {
@@ -33,7 +32,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -71,7 +70,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -143,7 +142,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -215,7 +214,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -243,7 +242,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -336,7 +335,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -471,7 +470,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 2"
           }
         ]
       },
@@ -500,7 +499,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 2"
           }
         ]
       },
@@ -664,7 +663,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -756,7 +755,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -785,7 +784,7 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
@@ -813,11 +812,53 @@
         "metrics": [
           {
             "type": "bandwidth",
-            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
       "attrs": { "K": 128, "N": 262144, "bits": 4, "block_size": 32 }
     }
   ]
 }

 {
   "tunableSpace": { "workgroupSize": [64, 128, 256] },
   "cases": [
     {
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.biasT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 2"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 2"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + numel(shapes.zeroPointsT) * 4 + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
         "metrics": [
           {
             "type": "bandwidth",
+            "value": "numel(shapes.bT) + (numel(shapes.aT) + numel(shapes.scalesT) + numel(shapes.yT)) * 4"
           }
         ]
       },
       "attrs": { "K": 128, "N": 262144, "bits": 4, "block_size": 32 }
+    },
+    {
+      "name": "mmnb-q4g32-m47-k4096-n4096-one-below-register-tile-floor-pathology",
+      "preset": "model",
+      "vars": { "M": 47, "K": 4096, "N": 4096, "bits": 4, "blockSize": 32 },
+      "inputs": {
+        "aT": { "shape": [47, 4096], "dtype": "float32", "dist": "normal", "seed": 643, "scale": 0.2 },
+        "bT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "q4pair", "seed": 644 },
+        "scalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 645,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [47, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.M * args.K * args.N" }] },
+      "attrs": { "K": 4096, "N": 4096, "bits": 4, "block_size": 32 }
+    },
+    {
+      "name": "mmnb-q4g32-m48-k4096-n4096-register-tile-floor-control",
+      "preset": "model",
+      "vars": { "M": 48, "K": 4096, "N": 4096, "bits": 4, "blockSize": 32 },
+      "inputs": {
+        "aT": { "shape": [48, 4096], "dtype": "float32", "dist": "normal", "seed": 646, "scale": 0.2 },
+        "bT": { "shape": [4096, 128, 16], "dtype": "uint8", "dist": "q4pair", "seed": 647 },
+        "scalesT": {
+          "shape": [4096, 128],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 648,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [48, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.M * args.K * args.N" }] },
+      "attrs": { "K": 4096, "N": 4096, "bits": 4, "block_size": 32 }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

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

build/webgpu/matmul-nbits-dp4a-quantize.wgsl.jinja CHANGED Viewed

@@ -13,11 +13,10 @@ var<workgroup> maxAbs: array<f32, WG>;
 @compute @workgroup_size(WG)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   // 2D-folded dispatch: wid.y carries the high bits past the per-dimension limit.
-  let wgFlat = wid.x + wid.y * nwg.x;
   let i = wgFlat * WG + lid.x;
   var v = vec4<f32>(0.0);
   if (i < VEC4_COUNT) {

 @compute @workgroup_size(WG)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   // 2D-folded dispatch: wid.y carries the high bits past the per-dimension limit.
+  let wgFlat = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let i = wgFlat * WG + lid.x;
   var v = vec4<f32>(0.0);
   if (i < VEC4_COUNT) {

build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja CHANGED Viewed

@@ -1,35 +1,43 @@
-{% if usesF16 %}
-enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
-// Each workgroup handles gemvNCols columns in groups of four, with one vec4
-// accumulator and partial array per group. Every column after the first is
-// bounds-guarded, so arbitrary N and group counts compose.
 {% set nCols = gemvNCols %}
 {% set groups = 2 if nCols == 8 else 1 %}
 {% set comps = ["x", "y", "z", "w"] %}
-{% set codesPerByte = 4 if bits == 2 else (2 if bits == 4 else 1) %}
-{% set codeMask = 3 if bits == 2 else (15 if bits == 4 else 255) %}
 const WG: u32 = {{ workgroupSize }}u;
 const N_COLS: u32 = {{ nCols }}u;
-const CODES: u32 = {{ codesPerByte }}u;
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 var<workgroup> partials{{ sfx }}: array<vec4<f32>, {{ workgroupSize }}>;
 {% endfor %}
-// {{ bits }}-bit unpack: q2 packs four codes per byte, q4 two nibbles, q8 one byte/weight.
-// The whole stored byte, addressed by its index within the block rather than by
-// one of the codes inside it.
-fn packed_byte(n: u32, block: u32, byte_index: u32) -> u32 {
-  return b[(n * params.kBlocks + block) * params.blobSize + byte_index];
-}
 // Per-(column, block) zero point. Symmetric weights use the baked default
-// ({{ defaultZero }}); asymmetric weights carry one widened u32 per (col, block)
-// in row-major [N, kBlocks] layout.
 fn zero_point({% if hasZero %}n: u32, block: u32{% endif %}) -> f32 {
 {% if hasZero %}
   return f32(zero_points[n * params.kBlocks + block]);
@@ -38,17 +46,27 @@ fn zero_point({% if hasZero %}n: u32, block: u32{% endif %}) -> f32 {
 {% endif %}
 }
 @compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>,
-  @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let tid = lid.x;
   // 2D-folded dispatch: wid.y carries the high bits past the per-dimension
   // limit (vocab-sized N). The fold over-dispatches, so the tail workgroups
   // return before any barrier.
-  let col_base = (wid.x + wid.y * nwg.x) * N_COLS;
   if (col_base >= params.N) {
     return;
   }
@@ -57,47 +75,74 @@ fn main(
   var acc{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
-  // Stride whole bytes: every code a byte holds is consumed from one fetch.
-  for (var k = tid * CODES; k < params.K; k = k + WG * CODES) {
-    let block = k / params.blockSize;
-    // blockSize is a power of two of at least 16, so a byte never straddles two
-    // blocks and every code below shares this block's scale and zero point.
-    let byte_index = (k % params.blockSize) / CODES;
-{% for j in range(codesPerByte) %}
-{% if j == 0 %}
-    let av0 = f32(a[k]);
 {% else %}
-    // K need not be a multiple of CODES; a code past the end contributes zero.
-    let av{{ j }} = select(0.0, f32(a[min(k + {{ j }}u, params.K - 1u)]), k + {{ j }}u < params.K);
-{% endif %}
 {% endfor %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}
 {% set i = g * 4 + c %}
 {% if i == 0 %}
-    let scale0 = f32(scales[(col_base + 0u) * params.kBlocks + block]);
-    let zero0 = zero_point({% if hasZero %}col_base + 0u, block{% endif %});
-    let word0 = packed_byte(col_base + 0u, block, byte_index);
-{% for j in range(codesPerByte) %}
-    acc{{ sfx }}.x = acc{{ sfx }}.x
-      + av{{ j }} * (f32((word0 >> {{ j * bits }}u) & {{ codeMask }}u) - zero0) * scale0;
-{% endfor %}
 {% else %}
     if (col_base + {{ i }}u < params.N) {
-      let scale{{ i }} = f32(scales[(col_base + {{ i }}u) * params.kBlocks + block]);
-      let zero{{ i }} = zero_point({% if hasZero %}col_base + {{ i }}u, block{% endif %});
-      let word{{ i }} = packed_byte(col_base + {{ i }}u, block, byte_index);
-{% for j in range(codesPerByte) %}
-      acc{{ sfx }}.{{ comps[c] }} = acc{{ sfx }}.{{ comps[c] }}
-        + av{{ j }} * (f32((word{{ i }} >> {{ j * bits }}u) & {{ codeMask }}u) - zero{{ i }}) * scale{{ i }};
-{% endfor %}
-    }
 {% endif %}
 {% endfor %}
 {% endfor %}
   }
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
   partials{{ sfx }}[tid] = acc{{ sfx }};
@@ -124,6 +169,7 @@ fn main(
 {% set sfx = "" if groups == 1 else g %}
     let total{{ sfx }} = partials{{ sfx }}[0];
 {% endfor %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}

+{% if useSubgroups %}
+enable subgroups;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
+// com.microsoft.MatMulNBits decode GEMV: one activation row against the packed
+// weight blob. A workgroup owns gemvNCols output columns; its lanes stride the
+// reduction axis in whole VECTORS of the blob -- a vec4<u32> of sixteen packed
+// bytes when the quantization block is a multiple of sixteen bytes, else one
+// u32 word -- so every weight byte arrives in the widest load the layout admits.
+// A vector never straddles a block, so one scale (and zero point) covers every
+// code it carries.
+//
+// Per vector and column the lane accumulates dot(a, codes) over the vector's
+// codes and applies the block's (scale, zero) once:
+//   sum_j a_j (c_j - z) s  ==  s (dot(a, c) - z sum_j a_j)
+// so the per-code work is a shift, a mask, a convert and one fma, and the
+// activation sum is shared by every column. Partials fold through workgroup
+// memory, so no subgroup feature is required.
 {% set nCols = gemvNCols %}
 {% set groups = 2 if nCols == 8 else 1 %}
 {% set comps = ["x", "y", "z", "w"] %}
 const WG: u32 = {{ workgroupSize }}u;
 const N_COLS: u32 = {{ nCols }}u;
+const VEC_PER_COL: u32 = {{ vecPerCol }}u;
+const VEC_PER_BLOCK: u32 = {{ vecPerBlock }}u;
+{% if actVec4 %}
+const VEC4_PER_VEC: u32 = {{ codesPerVec4 }}u;
+{% else %}
+const CODES_PER_VEC: u32 = {{ codesPerVec }}u;
+{% endif %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 var<workgroup> partials{{ sfx }}: array<vec4<f32>, {{ workgroupSize }}>;
 {% endfor %}
 // Per-(column, block) zero point. Symmetric weights use the baked default
+// ({{ defaultZero }}); asymmetric weights carry one widened value per (col,
+// block) in row-major [N, kBlocks] layout.
 fn zero_point({% if hasZero %}n: u32, block: u32{% endif %}) -> f32 {
 {% if hasZero %}
   return f32(zero_points[n * params.kBlocks + block]);
 {% endif %}
 }
+{% macro vec_dot(words) %}
+{% for w in range(vecWords) %}
+{% set word = (words ~ "." ~ comps[w]) if vecWords == 4 else words %}
+{% for h in range(codesPerWord) %}
+      dot = dot + a{{ w * codesPerWord + h }} * f32(({{ word }} >> {{ h * bits }}u) & {{ codeMask }}u);
+{% endfor %}
+{% endfor %}
+{%- endmacro %}
 @compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(
   @builtin(workgroup_id) wid: 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;
   // 2D-folded dispatch: wid.y carries the high bits past the per-dimension
   // limit (vocab-sized N). The fold over-dispatches, so the tail workgroups
   // return before any barrier.
+  let col_base = (wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u) * N_COLS;
   if (col_base >= params.N) {
     return;
   }
   var acc{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
+  for (var v = tid; v < VEC_PER_COL; v = v + WG) {
+    let block = v / VEC_PER_BLOCK;
+{% if actVec4 %}
+    let abase = v * VEC4_PER_VEC;
+{% for q in range(codesPerVec4) %}
+    let av{{ q }} = vec4<f32>(a[abase + {{ q }}u]);
+{% endfor %}
+{% for q in range(codesPerVec4) %}
+{% for c in range(4) %}
+    let a{{ q * 4 + c }} = av{{ q }}.{{ comps[c] }};
+{% endfor %}
+{% endfor %}
 {% else %}
+    let k0 = v * CODES_PER_VEC;
+{% for j in range(codesPerVec) %}
+    let a{{ j }} = select(0.0, f32(a[min(k0 + {{ j }}u, params.K - 1u)]), k0 + {{ j }}u < params.K);
 {% endfor %}
+{% endif %}
+    let asum = {% for j in range(codesPerVec) %}{{ " + " if j > 0 else "" }}a{{ j }}{% endfor %};
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}
 {% set i = g * 4 + c %}
 {% if i == 0 %}
+    {
 {% else %}
     if (col_base + {{ i }}u < params.N) {
 {% endif %}
+      let n = col_base + {{ i }}u;
+      let words = b[n * VEC_PER_COL + v];
+      let scale = f32(scales[n * params.kBlocks + block]);
+      let zero = zero_point({% if hasZero %}n, block{% endif %});
+      var dot = 0.0;
+{{ vec_dot("words") }}
+      acc{{ sfx }}.{{ comps[c] }} = acc{{ sfx }}.{{ comps[c] }} + (dot - zero * asum) * scale;
+    }
 {% endfor %}
 {% endfor %}
   }
+{% if useSubgroups %}
+  // Subgroup fold: one collective per accumulator, then the WG / subgroup-size
+  // per-subgroup partials fold once through workgroup memory. The lanes of a
+  // subgroup are contiguous in local_invocation_id, so tid / sgSize names it.
+  // Every lane reaches this point (the only early return is workgroup-uniform),
+  // so the collective runs in uniform control flow.
+{% for g in range(groups) %}
+{% set sfx = "" if groups == 1 else g %}
+  let sgSum{{ sfx }} = subgroupAdd(acc{{ sfx }});
+{% endfor %}
+  if (sgLane == 0u) {
+{% for g in range(groups) %}
+{% set sfx = "" if groups == 1 else g %}
+    partials{{ sfx }}[tid / sgSize] = sgSum{{ sfx }};
+{% endfor %}
+  }
+  workgroupBarrier();
+  if (tid == 0u) {
+    let subgroupCount = WG / sgSize;
+{% for g in range(groups) %}
+{% set sfx = "" if groups == 1 else g %}
+    var total{{ sfx }} = partials{{ sfx }}[0];
+    for (var i = 1u; i < subgroupCount; i = i + 1u) {
+      total{{ sfx }} = total{{ sfx }} + partials{{ sfx }}[i];
+    }
+{% endfor %}
+{% else %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
   partials{{ sfx }}[tid] = acc{{ sfx }};
 {% set sfx = "" if groups == 1 else g %}
     let total{{ sfx }} = partials{{ sfx }}[0];
 {% endfor %}
+{% endif %}
 {% for g in range(groups) %}
 {% set sfx = "" if groups == 1 else g %}
 {% for c in range(4) %}

build/webgpu/matmul-nbits-q4-dp4a-prefill.wgsl.jinja CHANGED Viewed

@@ -53,11 +53,14 @@ fn main(
       tA[r][w] = select(0u, a_quant[am * K4 + kt * 8u + w], am < M);
       let bn = nBase + r;
       if (bn < N) {
-        // Two blob bytes hold this word's four sequential weights
-        // (low nibble = even k). Rebias by the implicit zero point 8.
         let byteBase = (bn * KBLOCKS + block) * BLOB_SIZE + ((kBase % BLOCK_SIZE) + w * 4u) / 2u;
-        let b0 = b[byteBase];
-        let b1 = b[byteBase + 1u];
         let quads = vec4<i32>(i32(b0 & 15u), i32((b0 >> 4u) & 15u), i32(b1 & 15u), i32((b1 >> 4u) & 15u));
         tB[r][w] = pack4xI8(quads - vec4<i32>(8));
       } else {

       tA[r][w] = select(0u, a_quant[am * K4 + kt * 8u + w], am < M);
       let bn = nBase + r;
       if (bn < N) {
+        // Two blob bytes hold this word's four sequential weights (low nibble =
+        // even k). The blob is packed four bytes per u32 word, and an even byte
+        // index inside a 4-byte-aligned blob never straddles a word, so both
+        // bytes come from one load. Rebias by the implicit zero point 8.
         let byteBase = (bn * KBLOCKS + block) * BLOB_SIZE + ((kBase % BLOCK_SIZE) + w * 4u) / 2u;
+        let pair = b[byteBase >> 2u] >> ((byteBase & 3u) * 8u);
+        let b0 = pair & 255u;
+        let b1 = (pair >> 8u) & 255u;
         let quads = vec4<i32>(i32(b0 & 15u), i32((b0 >> 4u) & 15u), i32(b1 & 15u), i32((b1 >> 4u) & 15u));
         tB[r][w] = pack4xI8(quads - vec4<i32>(8));
       } else {

build/webgpu/matmul-nbits-q4-prefill-tile4x4.wgsl.jinja CHANGED Viewed

@@ -1,38 +1,28 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-const WG: u32 = {{ workgroupSize }}u;
 const TILE_M: u32 = 4u;
 const TILE_N: u32 = 4u;
 var<workgroup> partial0: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial1: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial2: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial3: array<vec4<f32>, {{ workgroupSize }}>;
-// {{ bits }}-bit unpack: q2 stores four codes per byte, q4 two, and q8 one.
-{% 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 %}
-{{- matmul_nbits_packed_code(bits=bits) }}
 fn store_row(row: u32, col_base: u32, value: vec4<f32>) {
   if (row >= params.rows) {
     return;
@@ -54,78 +44,103 @@ fn main(
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let tid = lid.x;
   // Column tiles fold into z past the per-dimension limit (vocab-sized N,
   // where y already carries the row-tile grid-stride). Over-dispatched tail
-  // workgroups return before any barrier.
-  let col_base = (wid.x + wid.z * nwg.x) * TILE_N;
-  if (col_base >= params.N) {
     return;
   }
   // Grid-stride over row tiles because dispatch y is device-clamped. This loop
   // reduces to one iteration when the dispatch covers every tile directly.
   for (var row_tile = wid.y; row_tile * TILE_M < params.rows; row_tile = row_tile + nwg.y) {
   let row_base = row_tile * TILE_M;
   var acc0 = vec4<f32>(0.0);
   var acc1 = vec4<f32>(0.0);
   var acc2 = vec4<f32>(0.0);
   var acc3 = vec4<f32>(0.0);
-  for (var block: u32 = 0u; block < params.kBlocks; block = block + 1u) {
-    let scale = vec4<f32>(
-      f32(scales[(col_base + 0u) * params.kBlocks + block]),
-      f32(scales[(col_base + 1u) * params.kBlocks + block]),
-      f32(scales[(col_base + 2u) * params.kBlocks + block]),
-      f32(scales[(col_base + 3u) * params.kBlocks + block])
-    );
-    let block_k = block * params.blockSize;
-    for (var offset: u32 = tid; offset < params.blockSize; offset = offset + WG) {
-      let k = block_k + offset;
 {% if hasZero %}
-      let zero = vec4<f32>(
-        f32(zero_points[(col_base + 0u) * params.kBlocks + block]),
-        f32(zero_points[(col_base + 1u) * params.kBlocks + block]),
-        f32(zero_points[(col_base + 2u) * params.kBlocks + block]),
-        f32(zero_points[(col_base + 3u) * params.kBlocks + block])
-      );
-{% else %}
-      let zero = vec4<f32>({{ defaultZero }});
 {% endif %}
-      let w = (vec4<f32>(
-        f32(packed_weight(col_base + 0u, block, offset)),
-        f32(packed_weight(col_base + 1u, block, offset)),
-        f32(packed_weight(col_base + 2u, block, offset)),
-        f32(packed_weight(col_base + 3u, block, offset))
-      ) - zero) * scale;
-      var a0 = 0.0;
-      var a1 = 0.0;
-      var a2 = 0.0;
-      var a3 = 0.0;
-      if (row_base + 0u < params.rows) {
-        a0 = f32(a[(row_base + 0u) * params.K + k]);
-      }
-      if (row_base + 1u < params.rows) {
-        a1 = f32(a[(row_base + 1u) * params.K + k]);
       }
-      if (row_base + 2u < params.rows) {
-        a2 = f32(a[(row_base + 2u) * params.K + k]);
       }
-      if (row_base + 3u < params.rows) {
-        a3 = f32(a[(row_base + 3u) * params.K + k]);
       }
-      acc0 = acc0 + a0 * w;
-      acc1 = acc1 + a1 * w;
-      acc2 = acc2 + a2 * w;
-      acc3 = acc3 + a3 * w;
     }
   }
   // Sync before reusing the shared partial buffers: on a subsequent grid-stride
-  // iteration the previous iteration's tid==0 store may still be reading them.
   workgroupBarrier();
   partial0[tid] = acc0;
   partial1[tid] = acc1;
@@ -133,38 +148,35 @@ fn main(
   partial3[tid] = acc3;
   workgroupBarrier();
-  var stride = WG / 2u;
-  loop {
-    if (stride == 0u) {
-      break;
-    }
-    if (tid < stride) {
       partial0[tid] = partial0[tid] + partial0[tid + stride];
       partial1[tid] = partial1[tid] + partial1[tid + stride];
       partial2[tid] = partial2[tid] + partial2[tid + stride];
       partial3[tid] = partial3[tid] + partial3[tid + stride];
     }
-    stride = stride / 2u;
     workgroupBarrier();
   }
-  if (tid == 0u) {
 {% if hasBias %}
     let biasv = vec4<f32>(
       f32(bias[col_base + 0u]),
-      f32(bias[col_base + 1u]),
-      f32(bias[col_base + 2u]),
-      f32(bias[col_base + 3u])
     );
-    store_row(row_base + 0u, col_base, partial0[0] + biasv);
-    store_row(row_base + 1u, col_base, partial1[0] + biasv);
-    store_row(row_base + 2u, col_base, partial2[0] + biasv);
-    store_row(row_base + 3u, col_base, partial3[0] + biasv);
 {% else %}
-    store_row(row_base + 0u, col_base, partial0[0]);
-    store_row(row_base + 1u, col_base, partial1[0]);
-    store_row(row_base + 2u, col_base, partial2[0]);
-    store_row(row_base + 3u, col_base, partial3[0]);
 {% endif %}
   }
   }

 {{ env.wgsl.resourceDeclarations }}
+// com.microsoft.MatMulNBits small-M prefill over the packed blob: any M >= 2, K
+// a whole number of quantization blocks. A workgroup covers TILE_M = 4
+// activation rows by TILE_N * COL_GROUPS output columns. Its lanes split into
+// COL_GROUPS groups of K_LANES; a group owns four adjacent columns and its
+// lanes advance through the blob one packed u32 word (32 / bits codes) at a
+// time. Each lane retains the four columns' words while applying them to four
+// activation rows, and one (scale, zero) pair per column serves every code in
+// the word. A word never straddles a quantization block.
+// Partials fold inside each column group through workgroup memory (portable:
+// no subgroup feature). K_LANES shrinks with the words per column for short K.
 const TILE_M: u32 = 4u;
 const TILE_N: u32 = 4u;
+const K_LANES: u32 = {{ kLanes }}u;
+const COL_GROUPS: u32 = {{ colGroups }}u;
+const COLS_PER_WORKGROUP: u32 = TILE_N * COL_GROUPS;
+const WORDS_PER_COL: u32 = {{ wordsPerCol }}u;
+const WORDS_PER_BLOCK: u32 = {{ wordsPerBlock }}u;
+const CODES_PER_WORD: u32 = {{ codesPerWord }}u;
 var<workgroup> partial0: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial1: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial2: array<vec4<f32>, {{ workgroupSize }}>;
 var<workgroup> partial3: array<vec4<f32>, {{ workgroupSize }}>;
 fn store_row(row: u32, col_base: u32, value: vec4<f32>) {
   if (row >= params.rows) {
     return;
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   let tid = lid.x;
+  let group = tid / K_LANES;
+  let kl = tid % K_LANES;
   // Column tiles fold into z past the per-dimension limit (vocab-sized N,
   // where y already carries the row-tile grid-stride). Over-dispatched tail
+  // workgroups return before any barrier; a column GROUP past N inside a live
+  // workgroup keeps running on zeroed loads so the barriers stay uniform, and
+  // its store is skipped.
+  let workgroup_col = (wid.x + wid.z * {{ DISPATCH_FOLD_WIDTH }}u) * COLS_PER_WORKGROUP;
+  if (workgroup_col >= params.N) {
     return;
   }
+  let col_base = workgroup_col + group * TILE_N;
+  let group_valid = col_base < params.N;
+  let col1 = col_base + 1u < params.N;
+  let col2 = col_base + 2u < params.N;
+  let col3 = col_base + 3u < params.N;
   // Grid-stride over row tiles because dispatch y is device-clamped. This loop
   // reduces to one iteration when the dispatch covers every tile directly.
   for (var row_tile = wid.y; row_tile * TILE_M < params.rows; row_tile = row_tile + nwg.y) {
   let row_base = row_tile * TILE_M;
+  // Rows past the end clamp onto the last real row; their accumulators are
+  // finite and the store guard drops them.
+  let row1 = row_base + 1u < params.rows;
+  let row2 = row_base + 2u < params.rows;
+  let row3 = row_base + 3u < params.rows;
+  let abase0 = row_base * params.K;
+  let abase1 = min(row_base + 1u, params.rows - 1u) * params.K;
+  let abase2 = min(row_base + 2u, params.rows - 1u) * params.K;
+  let abase3 = min(row_base + 3u, params.rows - 1u) * params.K;
   var acc0 = vec4<f32>(0.0);
   var acc1 = vec4<f32>(0.0);
   var acc2 = vec4<f32>(0.0);
   var acc3 = vec4<f32>(0.0);
+  for (var w = kl; w < WORDS_PER_COL; w = w + K_LANES) {
+    let block = w / WORDS_PER_BLOCK;
+    let k0 = w * CODES_PER_WORD;
+    // One packed word per column, plus its block's scale and zero point. A
+    // column past N reads nothing and dequantizes to zero (scale 0).
+    var q0 = 0u;
+    var q1 = 0u;
+    var q2 = 0u;
+    var q3 = 0u;
+    var scale = vec4<f32>(0.0);
+    var zero = vec4<f32>({{ defaultZero }});
+    if (group_valid) {
+      q0 = b[(col_base + 0u) * WORDS_PER_COL + w];
+      scale.x = f32(scales[(col_base + 0u) * params.kBlocks + block]);
 {% if hasZero %}
+      zero.x = f32(zero_points[(col_base + 0u) * params.kBlocks + block]);
+{% endif %}
+      if (col1) {
+        q1 = b[(col_base + 1u) * WORDS_PER_COL + w];
+        scale.y = f32(scales[(col_base + 1u) * params.kBlocks + block]);
+{% if hasZero %}
+        zero.y = f32(zero_points[(col_base + 1u) * params.kBlocks + block]);
 {% endif %}
       }
+      if (col2) {
+        q2 = b[(col_base + 2u) * WORDS_PER_COL + w];
+        scale.z = f32(scales[(col_base + 2u) * params.kBlocks + block]);
+{% if hasZero %}
+        zero.z = f32(zero_points[(col_base + 2u) * params.kBlocks + block]);
+{% endif %}
       }
+      if (col3) {
+        q3 = b[(col_base + 3u) * WORDS_PER_COL + w];
+        scale.w = f32(scales[(col_base + 3u) * params.kBlocks + block]);
+{% if hasZero %}
+        zero.w = f32(zero_points[(col_base + 3u) * params.kBlocks + block]);
+{% endif %}
       }
     }
+{% for j in range(codesPerWord) %}
+    {
+      let wt = (vec4<f32>(
+        f32((q0 >> {{ j * bits }}u) & {{ codeMask }}u),
+        f32((q1 >> {{ j * bits }}u) & {{ codeMask }}u),
+        f32((q2 >> {{ j * bits }}u) & {{ codeMask }}u),
+        f32((q3 >> {{ j * bits }}u) & {{ codeMask }}u)) - zero) * scale;
+      let k = k0 + {{ j }}u;
+      let a0 = f32(a[abase0 + k]);
+      let a1 = select(0.0, f32(a[abase1 + k]), row1);
+      let a2 = select(0.0, f32(a[abase2 + k]), row2);
+      let a3 = select(0.0, f32(a[abase3 + k]), row3);
+      acc0 = acc0 + a0 * wt;
+      acc1 = acc1 + a1 * wt;
+      acc2 = acc2 + a2 * wt;
+      acc3 = acc3 + a3 * wt;
+    }
+{% endfor %}
   }
   // Sync before reusing the shared partial buffers: on a subsequent grid-stride
+  // iteration the previous iteration's group-leader store may still be reading.
   workgroupBarrier();
   partial0[tid] = acc0;
   partial1[tid] = acc1;
   partial3[tid] = acc3;
   workgroupBarrier();
+  // Fold within each column group. K_LANES is a power of two, so kl + stride
+  // never leaves the group.
+  for (var stride = K_LANES / 2u; stride > 0u; stride = stride / 2u) {
+    if (kl < stride) {
       partial0[tid] = partial0[tid] + partial0[tid + stride];
       partial1[tid] = partial1[tid] + partial1[tid + stride];
       partial2[tid] = partial2[tid] + partial2[tid + stride];
       partial3[tid] = partial3[tid] + partial3[tid + stride];
     }
     workgroupBarrier();
   }
+  if (kl == 0u && group_valid) {
 {% if hasBias %}
     let biasv = vec4<f32>(
       f32(bias[col_base + 0u]),
+      select(0.0, f32(bias[min(col_base + 1u, params.N - 1u)]), col1),
+      select(0.0, f32(bias[min(col_base + 2u, params.N - 1u)]), col2),
+      select(0.0, f32(bias[min(col_base + 3u, params.N - 1u)]), col3)
     );
+    store_row(row_base + 0u, col_base, partial0[tid] + biasv);
+    store_row(row_base + 1u, col_base, partial1[tid] + biasv);
+    store_row(row_base + 2u, col_base, partial2[tid] + biasv);
+    store_row(row_base + 3u, col_base, partial3[tid] + biasv);
 {% else %}
+    store_row(row_base + 0u, col_base, partial0[tid]);
+    store_row(row_base + 1u, col_base, partial1[tid]);
+    store_row(row_base + 2u, col_base, partial2[tid]);
+    store_row(row_base + 3u, col_base, partial3[tid]);
 {% endif %}
   }
   }

build/webgpu/matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier — register-blocked.
@@ -30,6 +27,13 @@ const TN: u32 = {{ threadCols }}u;
 const WG_X: u32 = 16u;
 const WG_Y: u32 = 16u;
 const WG_THREADS: u32 = WG_X * WG_Y;
 // When A uses f16 storage, staging tileA at f16 is bit-identical to widening it
 // here. tileB stays f32 because the dequantized (code - zero)*scale product
@@ -46,18 +50,16 @@ var<workgroup> tileB: array<array<vec4<f32>, K_VECS>, BN>;   // [BN][BK/4]
 {% 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 %}
@@ -82,7 +84,13 @@ fn main(
 {% endfor %}
   let numTiles = (K + BK - 1u) / BK;
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
     let kBase = kt * BK;
     // A and B may have different tile areas (32x64 portable geometry), so stage
     // them independently. Every configured area is an exact WG_THREADS multiple.
@@ -92,17 +100,26 @@ fn main(
       let am = mBase + ar;
       let ak = kBase + ac4 * 4u;
       var aWord = vec4<{{ aTile }}>({{ aTile }}(0.0));
       if (am < M) {
 {% for component in range(4) %}
         if (ak + {{ component }}u < K) { aWord[{{ component }}u] = {{ aTile }}(a[am * K + ak + {{ component }}u]); }
 {% endfor %}
       }
       tileA[ar][ac4] = aWord;
     }
 {% if alignedBlockLoads %}
     // Every lane materializes one vector word of one output column. BLOCK_SIZE is
     // a multiple of BK, so the whole tile slice of a column shares one scale and
-    // zero point and each stored byte is read once for the codes it carries.
     for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + WG_THREADS) {
       let bc = idx / K_VECS;
       let kv = idx % K_VECS;
@@ -116,9 +133,9 @@ fn main(
         let blob_base = (bn * KBLOCKS + block) * BLOB_SIZE;
         let kk = kv * 4u;
 {% if bits == 2 %}
-        // Four 2-bit codes share a stored byte, and an aligned tile starts a
-        // word on a byte boundary, so one byte fills one vector word.
-        let packed = b[blob_base + (offset0 + kk) / 4u];
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
           word[{{ component }}u] = (f32((packed >> {{ component * 2 }}u) & 3u) - zero) * scale;
@@ -126,19 +143,19 @@ fn main(
 {% endfor %}
 {% elif bits == 4 %}
         // Two bytes carry the four K-adjacent codes of one vector word.
-        let byte_index = blob_base + (offset0 + kk) / 2u;
-        let packed0 = b[byte_index];
-        let packed1 = b[byte_index + 1u];
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
-          word[{{ component }}u] = (f32(({{ "packed0" if component < 2 else "packed1" }} >> {{ (component % 2) * 4 }}u) & 15u) - zero) * scale;
         }
 {% endfor %}
 {% else %}
         let byte_index = blob_base + offset0 + kk;
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
-          word[{{ component }}u] = (f32(b[byte_index + {{ component }}u] & 255u) - zero) * scale;
         }
 {% endfor %}
 {% endif %}
@@ -192,7 +209,11 @@ fn main(
   if (m0 + {{ i }}u < M) {
 {% for j in range(threadCols) %}
     if (n0 + {{ j }}u < N) {
       y[(m0 + {{ i }}u) * N + n0 + {{ j }}u] = {{ outputScalar }}(acc{{ i }}{{ j }}{% if hasBias %} + f32(bias[n0 + {{ j }}u]){% endif %});
     }
 {% endfor %}
   }

 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier — register-blocked.
 const WG_X: u32 = 16u;
 const WG_Y: u32 = 16u;
 const WG_THREADS: u32 = WG_X * WG_Y;
+{% set splitKValue = splitK if splitK is defined else 1 %}
+{% set tilesPerSplitValue = tilesPerSplit if tilesPerSplit is defined else 0 %}
+{% if splitKValue > 1 %}
+// Split-K: dispatch.z owns SPLIT_TILES consecutive K tiles and writes an f32
+// partial block for the combine pass (bias, if any, is added there).
+const SPLIT_TILES: u32 = {{ tilesPerSplitValue }}u;
+{% endif %}
 // When A uses f16 storage, staging tileA at f16 is bit-identical to widening it
 // here. tileB stays f32 because the dequantized (code - zero)*scale product
 {% 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 %}
 {% endfor %}
   let numTiles = (K + BK - 1u) / BK;
+{% if splitKValue > 1 %}
+  let ktBegin = wg.z * SPLIT_TILES;
+  let ktEnd = min(numTiles, ktBegin + SPLIT_TILES);
+  for (var kt: u32 = ktBegin; kt < ktEnd; kt = kt + 1u) {
+{% else %}
   for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+{% endif %}
     let kBase = kt * BK;
     // A and B may have different tile areas (32x64 portable geometry), so stage
     // them independently. Every configured area is an exact WG_THREADS multiple.
       let am = mBase + ar;
       let ak = kBase + ac4 * 4u;
       var aWord = vec4<{{ aTile }}>({{ aTile }}(0.0));
+{% if aVec4Loads is defined and aVec4Loads %}
+      // K is a multiple of four, so every staged word is one aligned vector load.
+      if (am < M) {
+        aWord = vec4<{{ aTile }}>(a[(am * K + ak) / 4u]);
+      }
+{% else %}
       if (am < M) {
 {% for component in range(4) %}
         if (ak + {{ component }}u < K) { aWord[{{ component }}u] = {{ aTile }}(a[am * K + ak + {{ component }}u]); }
 {% endfor %}
       }
+{% endif %}
       tileA[ar][ac4] = aWord;
     }
 {% if alignedBlockLoads %}
     // Every lane materializes one vector word of one output column. BLOCK_SIZE is
     // a multiple of BK, so the whole tile slice of a column shares one scale and
+    // zero point. The blob is packed four bytes per u32 word and the four
+    // K-adjacent codes of a vector word are byte-contiguous from a 4-aligned
+    // offset, so they always come out of ONE packed word: one load per vector.
     for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + WG_THREADS) {
       let bc = idx / K_VECS;
       let kv = idx % K_VECS;
         let blob_base = (bn * KBLOCKS + block) * BLOB_SIZE;
         let kk = kv * 4u;
 {% if bits == 2 %}
+        // Four 2-bit codes share a stored byte.
+        let byte_index = blob_base + ((offset0 + kk) >> 2u);
+        let packed = b[byte_index >> 2u] >> ((byte_index & 3u) * 8u);
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
           word[{{ component }}u] = (f32((packed >> {{ component * 2 }}u) & 3u) - zero) * scale;
 {% endfor %}
 {% elif bits == 4 %}
         // Two bytes carry the four K-adjacent codes of one vector word.
+        let byte_index = blob_base + ((offset0 + kk) >> 1u);
+        let packed = b[byte_index >> 2u] >> ((byte_index & 3u) * 8u);
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
+          word[{{ component }}u] = (f32((packed >> {{ component * 4 }}u) & 15u) - zero) * scale;
         }
 {% endfor %}
 {% else %}
         let byte_index = blob_base + offset0 + kk;
+        let packed = b[byte_index >> 2u];
 {% for component in range(4) %}
         if (kBase + kk + {{ component }}u < K) {
+          word[{{ component }}u] = (f32((packed >> {{ component * 8 }}u) & 255u) - zero) * scale;
         }
 {% endfor %}
 {% endif %}
   if (m0 + {{ i }}u < M) {
 {% for j in range(threadCols) %}
     if (n0 + {{ j }}u < N) {
+{% if splitKValue > 1 %}
+      y[wg.z * M * N + (m0 + {{ i }}u) * N + n0 + {{ j }}u] = acc{{ i }}{{ j }};
+{% else %}
       y[(m0 + {{ i }}u) * N + n0 + {{ j }}u] = {{ outputScalar }}(acc{{ i }}{{ j }}{% if hasBias %} + f32(bias[n0 + {{ j }}u]){% endif %});
+{% endif %}
     }
 {% endfor %}
   }

build/webgpu/matmul-nbits-q4-prefill-tiled.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier:
@@ -8,9 +5,9 @@ enable f16;
 // Register-blocked 32x32 shared-memory tiled GEMM (16x16 threads, each computes a
 // 2x2 micro-tile, K stepped in BK=16 chunks). The B operand is dequantized from
 // the packed blob straight into the shared B tile (no f32 weight matrix stored).
-// Each staged operand is reused across the micro-tile, and f32 accumulation is
-// exact for the 4/8-bit dequantized weight. Fully bounds-checked, so any M/N is
-// supported when K % blockSize == 0.
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
@@ -34,18 +31,16 @@ var<workgroup> tileB: array<array<f32, 32>, 16>;
 {% 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 %}

 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier:
 // Register-blocked 32x32 shared-memory tiled GEMM (16x16 threads, each computes a
 // 2x2 micro-tile, K stepped in BK=16 chunks). The B operand is dequantized from
 // the packed blob straight into the shared B tile (no f32 weight matrix stored).
+// Each staged operand is reused across the micro-tile. Dequantized weights and
+// accumulations use f32. Bounds checks support any M and N when K is divisible
+// by blockSize.
 const M: u32 = {{ M }}u;
 const K: u32 = {{ K }}u;
 const N: u32 = {{ N }}u;
 {% 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 %}

build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja CHANGED Viewed

@@ -2,9 +2,8 @@
 //   Y[M,N] = A[M,K] @ W[K,N] + bias?,  W[k,n] = (q4(n,block,offset) - zero) * scale
 // Uses a 32x64 or 64x64 subgroup-matrix output tile with an M-tail row_limit guard; the
 // B operand is materialized on the fly by dequantizing the packed q4 blob into
-// the shared B tile, so no f32 weight matrix is stored. f32 operands and f32
-// accumulation preserve precision; the dequantized weight is a 4-bit value, so
-// f32 is exact for it.
 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles; M tails handled by row_limit + the loadSHMA zero-fill.
 enable subgroups;
@@ -12,11 +11,9 @@ enable subgroups;
 enable subgroup_size_control;
 {% endif %}
 enable chromium_experimental_subgroup_matrix;
-{% if usesF16 %}
-enable f16;
-{% endif %}
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 const M: u32 = {{ M }}u;
@@ -68,25 +65,31 @@ fn loadSHMB(tile_base: u32, k_idx: u32, row: u32, c_idx: u32) {
 {% else %}
     let zero = {{ defaultZero }};
 {% endif %}
 {% if bits == 2 %}
-    // Four 2-bit codes share a stored byte; one code per iteration needs no assumption
-    // about where offset0 lands inside a byte.
     for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
       let code_offset = offset0 + i;
-      let packed_index = (n * KBLOCKS + block) * BLOB_SIZE + code_offset / 4u;
-      tile_B[row * TILE_K + col + i] = (f32((b[packed_index] >> ((code_offset % 4u) * 2u)) & 3u) - zero) * scale;
     }
 {% elif bits == 4 %}
-    for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 2u) {
-      let packed_index = (n * KBLOCKS + block) * BLOB_SIZE + (offset0 + i) / 2u;
-      let packed = b[packed_index];
-      tile_B[row * TILE_K + col + i] = (f32(packed & 15u) - zero) * scale;
-      tile_B[row * TILE_K + col + i + 1u] = (f32((packed >> 4u) & 15u) - zero) * scale;
     }
 {% else %}
-    for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
-      let packed_index = (n * KBLOCKS + block) * BLOB_SIZE + offset0 + i;
-      tile_B[row * TILE_K + col + i] = (f32(b[packed_index] & 255u) - zero) * scale;
     }
 {% endif %}
   } else {
@@ -147,15 +150,16 @@ fn main(
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
-{% set operandScalar = "f32" %}      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
-      var matA0: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset, TILE_K);
-      var matA1: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset + 8u * TILE_K, TILE_K);
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
-      var matB0: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset, TILE_K);
-      var matB1: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 8u * TILE_K, TILE_K);
-      var matB2: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 16u * TILE_K, TILE_K);
-      var matB3: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + 24u * TILE_K, TILE_K);
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

 //   Y[M,N] = A[M,K] @ W[K,N] + bias?,  W[k,n] = (q4(n,block,offset) - zero) * scale
 // Uses a 32x64 or 64x64 subgroup-matrix output tile with an M-tail row_limit guard; the
 // B operand is materialized on the fly by dequantizing the packed q4 blob into
+// the shared B tile, so no f32 weight matrix is stored. Dequantized operands and
+// accumulations use f32.
 // Alignment gates keep K % 32 == 0 and N % 64 == 0 so subgroupMatrixLoad never
 // sees partial 8x8 tiles; M tails handled by row_limit + the loadSHMA zero-fill.
 enable subgroups;
 enable subgroup_size_control;
 {% endif %}
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 const M: u32 = {{ M }}u;
 {% else %}
     let zero = {{ defaultZero }};
 {% endif %}
+    // The blob is packed four bytes per u32 word. The slice starts on an 8-code
+    // boundary inside a 4-byte-aligned blob, so every packed word below holds
+    // 32 / bits consecutive codes of this column and is read once.
+    let blob_base = (n * KBLOCKS + block) * BLOB_SIZE;
 {% if bits == 2 %}
+    // Sixteen 2-bit codes per word; the slice may start mid-word.
     for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 1u) {
       let code_offset = offset0 + i;
+      let byte_index = blob_base + (code_offset >> 2u);
+      let shift = (byte_index & 3u) * 8u + (code_offset & 3u) * 2u;
+      tile_B[row * TILE_K + col + i] = (f32((b[byte_index >> 2u] >> shift) & 3u) - zero) * scale;
     }
 {% elif bits == 4 %}
+    for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 8u) {
+      let packed = b[(blob_base + ((offset0 + i) >> 1u)) >> 2u];
+{% for h in range(8) %}
+      tile_B[row * TILE_K + col + i + {{ h }}u] = (f32((packed >> {{ h * 4 }}u) & 15u) - zero) * scale;
+{% endfor %}
     }
 {% else %}
+    for (var i = 0u; i < {{ bLoadWidth }}u; i = i + 4u) {
+      let packed = b[(blob_base + offset0 + i) >> 2u];
+{% for h in range(4) %}
+      tile_B[row * TILE_K + col + i + {{ h }}u] = (f32((packed >> {{ h * 8 }}u) & 255u) - zero) * scale;
+{% endfor %}
     }
 {% endif %}
   } else {
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
+      let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
+{% for r in range(2) %}
+      var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});
+{% endfor %}
       let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for c in range(4) %}
+      var matB{{ c }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&{{ "xm" if directInputs else "tile_B" }}, matrix_b_offset{% if c > 0 %} + {{ c * 8 }}u * TILE_K{% endif %}, {{ "N" if directInputs else "TILE_K" }});
+{% endfor %}
       matC00 = subgroupMatrixMultiplyAccumulate(matA0, matB0, matC00);
       matC01 = subgroupMatrixMultiplyAccumulate(matA0, matB1, matC01);

build/webgpu/matmul-nbits.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -12,29 +9,27 @@ const WG: u32 = {{ workgroupSize }}u;
 {% 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 %}
 {{- matmul_nbits_packed_code(bits=bits) }}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat output-element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit. When nwg.y == 1 this
   // reduces to gid.x; the index >= total guard drops the tail.
-  let index = gid.x + gid.y * nwg.x * WG;
   let total = params.rows * params.N;
   if (index >= total) {

 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
 {% 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 %}
 {{- matmul_nbits_packed_code(bits=bits) }}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat output-element index: gid.y carries the high bits past the
+  // per-axis dispatch fold width. With no fold this
   // reduces to gid.x; the index >= total guard drops the tail.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.rows * params.N;
   if (index >= total) {

build/webgpu/metadata.json CHANGED Viewed

@@ -1,25 +1,63 @@
 {
   "name": "com.microsoft.MatMulNBits",
-  "id": "_com_microsoft_matmulnbits_webgpu_5495bd2",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "BbpRtHBH/vkJDaybyRBNKhuUCpe4XMJtKmW3bhRcEZE=",
-      "manifest.json": "z8oCYWRaMHew5ogYuf7Txv4/6y2qpVoVODDlLxEZ+vA=",
-      "matmul-nbits-dp4a-quantize.wgsl.jinja": "WEUGSDH5qWYg10f/mkAHxq8AkIkMW9R/Wn/Y23Q0Jw0=",
-      "matmul-nbits-gemv-q4.wgsl.jinja": "x1LxVx2cJBdN2cqapXz2T6AScpTcbmTO+B3tzvIXuT8=",
-      "matmul-nbits-q4-dp4a-prefill.wgsl.jinja": "qUrURZDpPzMpABnTgESkHgLCQiAWxEZ8wvzkAPTW2kI=",
-      "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja": "HjZraSjOpaSl5Cjns4scselOIn+ibxFEu+9LTeUI0QI=",
-      "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja": "oU7BcxJnKIEfgnkWHPq77RrCmrE0iIm6ED7cgVfZUYs=",
-      "matmul-nbits-q4-prefill-tiled.wgsl.jinja": "TdetCiv5o88hvWCGNDewpFX+M+NZrz9rekU/d5kMZTM=",
-      "matmul-nbits-q4-sgmat.wgsl.jinja": "PfO9QrUyNflp7RtUbUTRDwx7GDI6UUwz9I+6mJpO1OA=",
-      "matmul-nbits.wgsl.jinja": "DRkwODEER11pYQxt2ssS9YvFggO3CbUqQ9yag53J5lY=",
-      "test.json": "XLwnbLdOJASke2yBnDScjJi+o2XISmdRCebb2jwWpDI="
     }
   },
-  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulNBits" }
 }

 {
   "name": "com.microsoft.MatMulNBits",
+  "id": "_com_microsoft_matmulnbits_webgpu_6f18c00",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "P8wNyT1cbaSLn1VS4hVRFQxLYwji4Z7gp69Mt2pQaFo=",
+      "manifest.json": "QPPdzEGo1lJtp16tQaaqA9UHKPlvBY3fq66gG4XPvyQ=",
+      "matmul-nbits-dp4a-quantize.wgsl.jinja": "0gqEvgBRzH2Demz/RvrFkUOV487GyCq7Ujd26GyD+xw=",
+      "matmul-nbits-gemv-q4.wgsl.jinja": "ewlLPcW7t3UdnoLymZBoXnV+1oOeB9+YTyaOAnPhhfA=",
+      "matmul-nbits-q4-dp4a-prefill.wgsl.jinja": "pd9dWWUdIgYFVCRayU5OlvjYIvyqPvkOGgwtEPNMHas=",
+      "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja": "8DVy3szxcxwVItIlEXoioQi2YxDmmd5TCFW4BCYPiCQ=",
+      "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja": "W6JEeRleux7Jm3wmaflEusYFNuCiFtXnBTEW5cfS4Fk=",
+      "matmul-nbits-q4-prefill-tiled.wgsl.jinja": "eAEIbpkW0qziqyXhYNRiVbCuDRghjoz8EjnOrP8AgD8=",
+      "matmul-nbits-q4-sgmat.wgsl.jinja": "SFEMothi+irkTIclMjeTF1sCmB61Ai/2hGgnoOyNMUM=",
+      "matmul-nbits.wgsl.jinja": "UndxgqiV/O19Plpxda1d388lnQWyeCugbBvJGcYr9t8=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "Yz1hjK55R/kndUrw3ugPqgPaOBwKmYupqZoVdJTHO5Q=",
+      "test.json": "Hnb4z+ExdwYZb4EjSBR0gdGi5eY1kh7DaoGj3UTC5Ug="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "q4_dp4a_prefill": ["matmul-nbits-dp4a-quantize.wgsl.jinja", "matmul-nbits-q4-dp4a-prefill.wgsl.jinja"],
+      "gemv_default_zero": ["matmul-nbits-gemv-q4.wgsl.jinja"],
+      "prefill_sgmat_default_zero": ["matmul-nbits-q4-sgmat.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_splitk_default_zero": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_default_zero": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_reg_default_zero": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_default_zero": ["matmul-nbits-q4-prefill-tiled.wgsl.jinja"],
+      "prefill_tile4x4_default_zero": ["matmul-nbits-q4-prefill-tile4x4.wgsl.jinja"],
+      "default_zero": ["matmul-nbits.wgsl.jinja"],
+      "gemv_zero_bias": ["matmul-nbits-gemv-q4.wgsl.jinja"],
+      "prefill_sgmat_zero_bias": ["matmul-nbits-q4-sgmat.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_splitk_zero_bias": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_zero_bias": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_reg_zero_bias": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_zero_bias": ["matmul-nbits-q4-prefill-tiled.wgsl.jinja"],
+      "prefill_tile4x4_zero_bias": ["matmul-nbits-q4-prefill-tile4x4.wgsl.jinja"],
+      "zero_bias": ["matmul-nbits.wgsl.jinja"],
+      "gemv_zero_only": ["matmul-nbits-gemv-q4.wgsl.jinja"],
+      "prefill_sgmat_zero_only": ["matmul-nbits-q4-sgmat.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_splitk_zero_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_zero_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_reg_zero_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_zero_only": ["matmul-nbits-q4-prefill-tiled.wgsl.jinja"],
+      "prefill_tile4x4_zero_only": ["matmul-nbits-q4-prefill-tile4x4.wgsl.jinja"],
+      "zero_only": ["matmul-nbits.wgsl.jinja"],
+      "gemv_bias_only": ["matmul-nbits-gemv-q4.wgsl.jinja"],
+      "prefill_sgmat_bias_only": ["matmul-nbits-q4-sgmat.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_splitk_bias_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja", "reduce-axis0-splitk-combine.wgsl.jinja"],
+      "prefill_tiled_reg_vec4_bias_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_reg_bias_only": ["matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja"],
+      "prefill_tiled_bias_only": ["matmul-nbits-q4-prefill-tiled.wgsl.jinja"],
+      "prefill_tile4x4_bias_only": ["matmul-nbits-q4-prefill-tile4x4.wgsl.jinja"],
+      "bias_only": ["matmul-nbits.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,122 @@

+// Pass 2 of the split-K column-wise reduction. One invocation per output column
+// folds the segment partials and applies the selected reduction's final step.
+// Segments are folded in ascending order for deterministic results. This order
+// differs from the single-pass reduction but remains within the f32 tolerance.
+{% set addBias = addBias is defined and addBias %}
+{% set biasCols = biasCols | default(0) %}
+{% set intMode = intMode is defined and intMode %}
+{% set yv = "f16(" if outputF16 else "" %}
+{% set vy = ")" if outputF16 else "" %}
+{% if outputF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
+/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
+ * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
+{% set logicalBoolIdentity = logicalBool is defined and logicalBool %}
+fn {{ name }}() -> {{ scalar }} {
+{% if scalar == "i32" %}
+  return {{ "-2147483647i - 1i" if op == "max" else "2147483647i" }};
+{% elif scalar == "u32" %}
+  return {{ "1u" if logicalBoolIdentity and op == "min" else "0u" if op == "max" else "4294967295u" }};
+{% else %}
+  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
+  return bitcast<f32>(bits);
+{% endif %}
+}
+{%- endmacro %}
+const WG: u32 = {{ workgroupSize }}u;
+const SPLIT: u32 = {{ split }}u;
+{% if addBias %}
+const BIAS_COLS: u32 = {{ biasCols }}u;
+{% endif %}
+{% if op == "logsumexp" %}
+const F32_MIN: f32 = -3.4028234663852886e38;
+const F32_MAX: f32 = 3.4028234663852886e38;
+fn is_nan_f32(value: f32) -> bool {
+  let bits = bitcast<u32>(value);
+  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
+}
+{% elif op == "max" or op == "min" %}
+{{ wgsl_minmax_identity("reduction_identity", op) }}
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  let stride = nwg.x * WG;
+  let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x;
+  for (var col = start; col < params.cols; col = col + stride) {
+{% if op == "logsumexp" %}
+    // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers.
+    var nan_value = 0.0;
+    var has_nan = false;
+    var global_max = F32_MIN;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let nv = partials[(2u * SPLIT + seg) * params.cols + col];
+      if (nv != 0.0 || is_nan_f32(nv)) {
+        has_nan = true;
+        nan_value = nv;
+      }
+      global_max = max(global_max, partials[seg * params.cols + col]);
+    }
+    var sum = 0.0;
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let seg_max = partials[seg * params.cols + col];
+      let seg_sum = partials[(SPLIT + seg) * params.cols + col];
+      sum = sum + seg_sum * exp(seg_max - global_max);
+    }
+    let has_positive_inf = global_max > F32_MAX;
+    let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf);
+    y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
+{% else %}
+{% if intMode %}
+{% if op == "prod" %}
+    var total = 1i;
+{% else %}
+    var total = 0i;
+{% endif %}
+{% else %}
+{% if op == "max" %}
+    var total = reduction_identity();
+{% elif op == "min" %}
+    var total = reduction_identity();
+{% elif op == "prod" %}
+    var total = 1.0;
+{% else %}
+    var total = 0.0;
+{% endif %}
+{% endif %}
+    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
+      let p = partials[seg * params.cols + col];
+{% if op == "max" or op == "min" %}
+      total = {{ op }}(total, p);
+{% elif op == "prod" %}
+      total = total * p;
+{% else %}
+      total = total + p;
+{% endif %}
+    }
+{% if addBias %}
+    total = total + f32(bias[col % BIAS_COLS]);
+{% endif %}
+{% if op == "l2" %}
+    y[col] = {{ yv }}sqrt(total){{ vy }};
+{% elif op == "logsum" %}
+    y[col] = {{ yv }}log(total){{ vy }};
+{% elif op == "mean" %}
+    y[col] = {{ yv }}total / f32(params.rows){{ vy }};
+{% else %}
+{% if outputF16 %}
+    y[col] = f16(total);
+{% else %}
+    y[col] = total;
+{% endif %}
+{% endif %}
+{% endif %}
+  }
+}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.MatMulNBits",
   "fixtureArrays": {
     "q4_prefill_tiled_reg_tailk_m32_k33_n4096_input_bT": [16, 50, 84, 118, 152, 186, 220, 254, 135, 119, 102, 85, 68, 51, 34, 17],
     "q4_weight_cycle_b_t": [16, 50, 84, 118, 152, 186, 220, 254, 33, 67, 101, 135, 169, 203, 237, 15],
@@ -78,7 +77,7 @@
     {
       "name": "q8_zero_bias_prefill_tile4x4_small_m8",
       "provenance": {
-        "notes": "q8 companion to q4_zero_bias_prefill_tile4x4_small_m8: the relaxed tile4x4 bias floor (M>=2) also makes prefill_tile4x4_zero_bias selectable for q8 prefill+bias at M in [2,64). Validates the q8 unpack branch + bias epilogue on the tile4x4 path at small M."
       },
       "inputs": {
         "aT": {
@@ -181,7 +180,7 @@
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/matmul_4bits_test.cc",
         "test": "MatMulNBits.UnsupportedBlockSize_512",
-        "notes": "The public schema permits every power-of-two block size of at least 16, and this package implements block_size=512. ORT's installed, pinned, and inspected current CPU providers instead restrict execution to 16, 32, 64, 128, or 256. With A=1, packed q4 codes=9, default zero point=8, and scale=0.25, the exact result is 512 * (9 - 8) * 0.25 = 128."
       },
       "inputs": {
         "aT": { "dtype": "float32", "shape": [1, 512], "data": { "kind": "constant", "value": 1.0 } },
@@ -242,7 +241,9 @@
     },
     {
       "name": "q4_zero_bias_prefill_sgmat_m32_storage_floor",
-      "provenance": { "notes": "Zero-point/bias sibling of the 32-row subgroup-matrix storage-floor regression." },
       "inputs": {
         "aT": {
           "dtype": "float32",
@@ -276,7 +277,7 @@
     {
       "name": "q4_no_zero_prefill_g64_sgmat",
       "provenance": {
-        "notes": "Compact blockSize=64 q4 default-zero prefill sibling for model g64 benchmarks. Full tier routes through prefill_sgmat_default_zero; no-MMA/no-subgroups/baseline route through prefill_tiled_default_zero."
       },
       "inputs": {
         "aT": {
@@ -323,7 +324,7 @@
     {
       "name": "q8_no_zero_prefill_sgmat_m32_storage_floor",
       "provenance": {
-        "notes": "q8 sibling proving the 32-row subgroup-matrix tile remains available at the 16 KiB storage floor."
       },
       "inputs": {
         "aT": {
@@ -377,9 +378,131 @@
       "outputs": { "yT": { "dtype": "float32", "shape": [128, 64], "tolerance": 0.001 } },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
     {
       "name": "q8_zero_bias_prefill_sgmat_m32_storage_floor",
-      "provenance": { "notes": "q8 zero-point/bias sibling of the 32-row subgroup-matrix storage-floor regression." },
       "inputs": {
         "aT": {
           "dtype": "float32",
@@ -692,9 +815,7 @@
         }
       },
       "outputs": { "yT": { "dtype": "float16", "shape": [512, 1024], "tolerance": 0.0001, "relTolerance": 0.001 } },
-      "provenance": {
-        "notes": "f16 activations on the register-tiled prefill path: this is the only shape class where the shared A/B tiles are staged in f16, so it is the fixture that covers that staging."
-      },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
     {
@@ -772,7 +893,7 @@
       },
       "outputs": { "yT": { "dtype": "float16", "shape": [512, 1024], "tolerance": 0.0001, "relTolerance": 0.001 } },
       "provenance": {
-        "notes": "q8 sibling of q4_f16_prefill_reg_m512_n1024. 8-bit codes give the widest dequantized range, so the f16 tile store rounds hardest here."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
@@ -836,7 +957,7 @@
     {
       "name": "q4_gemv_default_zero_m1_tail_n7",
       "provenance": {
-        "notes": "M=1 GEMV with N=7 (not a multiple of 4). Exercises the vec4 tail-column store guards (col_base+1/2/3 < N) in matmul-nbits-gemv-q4.wgsl.jinja; existing M==1 tests all use N=4. Routes to gemv_default_zero (prio 20)."
       },
       "inputs": {
         "aT": {
@@ -861,7 +982,7 @@
     {
       "name": "q4_gemv_default_zero_m1_odd_k33_n5",
       "provenance": {
-        "notes": "K=33 is odd, so the last stored byte of the second block carries one live nibble and one past the end of K. The GEMV walks whole bytes, so this is the only shape that exercises its out-of-range activation guard; every other M==1 case has an even K."
       },
       "inputs": {
         "aT": {
@@ -886,7 +1007,7 @@
     {
       "name": "q2_gemv_default_zero_m1_k18_n5",
       "provenance": {
-        "notes": "The only two-bit M==1 case: four codes share one stored byte, and K=18 leaves two of the last byte's four codes past the end of K. Covers both the two-bit unpack in the GEMV and its out-of-range activation guard at the widest code-per-byte packing."
       },
       "inputs": {
         "aT": {
@@ -938,7 +1059,7 @@
       "name": "q4_gemv_default_zero_m1_tail_n7_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
-        "notes": "GEMV_N_COLS=8 twin of tail_n7: one workgroup, second column group entirely tail-guarded (cols 4-6 active, 7 out of range)."
       },
       "inputs": {
         "aT": {
@@ -963,7 +1084,7 @@
     {
       "name": "q8_zero_bias_gemv_m1_tail_n5",
       "provenance": {
-        "notes": "M=1 q8 GEMV with N=5 (not a multiple of 4), WITH zero_points and bias. Exercises the vec4 tail-column store guards together with the q8 unpack and bias epilogue in matmul-nbits-gemv-q4.wgsl.jinja. Routes to gemv_zero_bias (prio 20)."
       },
       "inputs": {
         "aT": {
@@ -1000,7 +1121,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "M=1 q8 GEMV with explicit zero points and no bias. This pins the zero-points-only binding layout and epilogue instead of relying on the combined zero-point-plus-bias case."
       },
       "inputs": {
         "aT": {
@@ -1032,7 +1153,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "M=1 q8 GEMV with bias and the schema-default zero point. This pins the bias-only binding layout and epilogue instead of relying on the combined zero-point-plus-bias case."
       },
       "inputs": {
         "aT": {
@@ -1063,7 +1184,7 @@
       "name": "q8_zero_bias_gemv_m1_tail_n5_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
-        "notes": "GEMV_N_COLS=8 twin of tail_n5: q8 unpack + zero_points + bias epilogue with only one live column in the second group."
       },
       "inputs": {
         "aT": {
@@ -1098,7 +1219,7 @@
     {
       "name": "q8_zero_bias_naive_fallback_tailK_m3_n6",
       "provenance": {
-        "notes": "K=17,blockSize=16 gives a tail block; N=6 is below the register-tiled floor and the aligned tile paths are ineligible, while M=3 bypasses GEMV. With zero_points+bias present this routes to the priority-0 zero_bias scalar fallback and covers q8 unpack, per-block zero points, bias, and a partial K block."
       },
       "inputs": {
         "aT": {
@@ -1161,8 +1282,8 @@
     {
       "name": "q4_prefill_tiled_reg_aligned_m32_k32_n4096_at_workgroup_floor",
       "provenance": {
-        "source": "authored for route coverage",
-        "notes": "Aligned-K twin of the tail-K case at the same M and N. A tail-K shape skips the register tile's workgroup floor entirely, so no fixture exercised that floor; this one sits exactly on it at 64 workgroups."
       },
       "inputs": {
         "aT": {
@@ -1190,7 +1311,7 @@
     {
       "name": "q4_prefill_tiled_reg_tailk_m32_k33_n512_below_wide_n_floor",
       "provenance": {
-        "notes": "Tail-block coverage below the ordinary wide register-tile N floor. K=33 exercises the padded final block, while N=512 keeps the case inside the relaxed band that the larger N=4096 case cannot distinguish."
       },
       "inputs": {
         "aT": {
@@ -1273,7 +1394,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1305,7 +1426,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1360,7 +1481,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
@@ -1392,7 +1513,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
@@ -1424,7 +1545,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
@@ -1456,7 +1577,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
@@ -1488,7 +1609,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
@@ -1520,7 +1641,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
@@ -1529,7 +1650,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1561,7 +1682,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1593,7 +1714,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1625,7 +1746,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "inputs": {
         "aT": {
@@ -1680,7 +1801,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 64, "bits": 4, "block_size": 32 }
     },
@@ -1712,7 +1833,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
-        "notes": "zero_points-only and bias-only inputs must stay on the fast kernels instead of falling to the naive per-element path."
       },
       "attrs": { "K": 128, "N": 64, "bits": 4, "block_size": 32 }
     },
@@ -1738,7 +1859,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc",
         "test": "accuracy_level=4 int8 activation quantization (DP4AMatMulQuantize + DP4AMatMulNBits)",
-        "notes": "accuracy_level 4 permits int8 activation math: A is quantized per 128-element block (scale = max_abs / 127, values pack4x8snorm-rounded) and multiplied against the 4-bit weights with packed integer dots. The reference stays float, so tolerances absorb the permitted quantization error."
       },
       "attrs": { "K": 256, "N": 64, "bits": 4, "block_size": 32, "accuracy_level": 4 }
     },
@@ -1753,7 +1874,7 @@
         "bT": {
           "dtype": "uint8",
           "shape": [16, 2, 32],
-          "data": { "kind": "cycle", "values": [37, 201, 116, 88, 250, 13, 164, 95, 142, 7, 233] }
         },
         "scalesT": { "dtype": "float32", "shape": [16, 2], "data": { "kind": "linspace", "start": 0.02, "end": 0.06 } }
       },
@@ -1761,7 +1882,7 @@
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc",
         "test": "accuracy_level=4 int8 activation quantization (DP4AMatMulQuantize + DP4AMatMulNBits)",
-        "notes": "accuracy_level 4 permits int8 activation math: A is quantized per 128-element block (scale = max_abs / 127, values pack4x8snorm-rounded) and multiplied against the 4-bit weights with packed integer dots. The reference stays float, so tolerances absorb the permitted quantization error."
       },
       "attrs": { "K": 128, "N": 16, "bits": 4, "block_size": 64, "accuracy_level": 4 }
     },
@@ -2012,6 +2133,151 @@
       },
       "outputs": { "yT": { "dtype": "float32", "shape": [3, 6], "tolerance": 0.001 } },
       "attrs": { "K": 17, "N": 6, "bits": 8, "block_size": 16 }
     }
   ]
 }

 {
   "fixtureArrays": {
     "q4_prefill_tiled_reg_tailk_m32_k33_n4096_input_bT": [16, 50, 84, 118, 152, 186, 220, 254, 135, 119, 102, 85, 68, 51, 34, 17],
     "q4_weight_cycle_b_t": [16, 50, 84, 118, 152, 186, 220, 254, 33, 67, 101, 135, 169, 203, 237, 15],
     {
       "name": "q8_zero_bias_prefill_tile4x4_small_m8",
       "provenance": {
+        "notes": "An eight-row q8 prefill with zero points and bias exercises q8 unpacking and the bias epilogue on the small-M tile4x4 route."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/test/contrib_ops/matmul_4bits_test.cc",
         "test": "MatMulNBits.UnsupportedBlockSize_512",
+        "notes": "The schema permits power-of-two block sizes of at least 16, including 512. With A=1, q4 code 9, default zero point 8, and scale 0.25, the result is exactly `512 * (9 - 8) * 0.25 = 128`."
       },
       "inputs": {
         "aT": { "dtype": "float32", "shape": [1, 512], "data": { "kind": "constant", "value": 1.0 } },
     },
     {
       "name": "q4_zero_bias_prefill_sgmat_m32_storage_floor",
+      "provenance": {
+        "notes": "Explicit zero points and bias exercise the 32-row subgroup-matrix tile at the 16 KiB workgroup-storage floor."
+      },
       "inputs": {
         "aT": {
           "dtype": "float32",
     {
       "name": "q4_no_zero_prefill_g64_sgmat",
       "provenance": {
+        "notes": "A compact q4 prefill with block size 64 and the default zero point exercises the subgroup-matrix route and its portable tiled fallback."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_no_zero_prefill_sgmat_m32_storage_floor",
       "provenance": {
+        "notes": "A q8 prefill exercises the 32-row subgroup-matrix tile at the 16 KiB workgroup-storage floor."
       },
       "inputs": {
         "aT": {
       "outputs": { "yT": { "dtype": "float32", "shape": [128, 64], "tolerance": 0.001 } },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
+    {
+      "name": "q8_no_zero_prefill_reg_vec4_splitk_m128_k1024_n1024",
+      "provenance": {
+        "notes": "Without zero points or bias, 32 output tiles over K=1024 use two K slices whose float32 partial sums are combined."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.35 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024, 32, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/mixed_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [128, 1024], "tolerance": 0.005 } },
+      "attrs": { "K": 1024, "N": 1024, "bits": 8, "block_size": 32 }
+    },
+    {
+      "name": "q8_zero_bias_prefill_reg_vec4_splitk_m128_k1024_n1024",
+      "provenance": {
+        "notes": "With zero points and bias, the split-K four-wide route applies zero points in each partial pass and adds bias once in the combine."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.35 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024, 32, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/mixed_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "zeroPointsT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": [120.0, 132.0, 128.0, 140.0, 135.0, 118.0, 144.0, 126.0] }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [128, 1024], "tolerance": 0.005 } },
+      "attrs": { "K": 1024, "N": 1024, "bits": 8, "block_size": 32 }
+    },
+    {
+      "name": "q8_zero_only_prefill_reg_vec4_splitk_m128_k1024_n1024",
+      "provenance": {
+        "notes": "With zero points and no bias, the split-K four-wide route applies zero points in each partial pass and only sums in the combine."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.35 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024, 32, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/mixed_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "zeroPointsT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": [120.0, 132.0, 128.0, 140.0, 135.0, 118.0, 144.0, 126.0] }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [128, 1024], "tolerance": 0.005 } },
+      "attrs": { "K": 1024, "N": 1024, "bits": 8, "block_size": 32 }
+    },
+    {
+      "name": "q8_bias_only_prefill_reg_vec4_splitk_m128_k1024_n1024",
+      "provenance": {
+        "notes": "With bias and the implicit zero point, the split-K four-wide route adds bias once after combining the partial sums."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.35 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024, 32, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/mixed_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [1024, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [128, 1024], "tolerance": 0.005 } },
+      "attrs": { "K": 1024, "N": 1024, "bits": 8, "block_size": 32 }
+    },
     {
       "name": "q8_zero_bias_prefill_sgmat_m32_storage_floor",
+      "provenance": {
+        "notes": "Q8 weights with explicit zero points and bias exercise the 32-row subgroup-matrix tile at the 16 KiB storage floor."
+      },
       "inputs": {
         "aT": {
           "dtype": "float32",
         }
       },
       "outputs": { "yT": { "dtype": "float16", "shape": [512, 1024], "tolerance": 0.0001, "relTolerance": 0.001 } },
+      "provenance": { "notes": "Float16 activations exercise float16 A/B staging on the register-tiled prefill path." },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
     {
       },
       "outputs": { "yT": { "dtype": "float16", "shape": [512, 1024], "tolerance": 0.0001, "relTolerance": 0.001 } },
       "provenance": {
+        "notes": "Eight-bit weights give the widest dequantized range, exercising float16 staging and store rounding on the register-tiled prefill path."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
     {
       "name": "q4_gemv_default_zero_m1_tail_n7",
       "provenance": {
+        "notes": "M=1 and N=7 leave three live columns in the final four-column output group. The guarded stores must write those columns without accessing the eighth column; zero points are omitted."
       },
       "inputs": {
         "aT": {
     {
       "name": "q4_gemv_default_zero_m1_odd_k33_n5",
       "provenance": {
+        "notes": "With K=33, the last stored byte of the second q4 block contains one live nibble and one beyond K. The GEMV must guard the out-of-range activation."
       },
       "inputs": {
         "aT": {
     {
       "name": "q2_gemv_default_zero_m1_k18_n5",
       "provenance": {
+        "notes": "Four q2 codes share each byte, and K=18 leaves two codes in the final byte beyond K. The GEMV must unpack the live codes and guard both out-of-range activations."
       },
       "inputs": {
         "aT": {
       "name": "q4_gemv_default_zero_m1_tail_n7_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
+        "notes": "With `GEMV_N_COLS = 8` and N=7, one workgroup has active columns 4 through 6 and a fully guarded column 7 in its second group."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_zero_bias_gemv_m1_tail_n5",
       "provenance": {
+        "notes": "M=1 and N=5 leave one live column in the final four-column output group. The case combines guarded tail stores with q8 unpacking, explicit zero points, and the bias epilogue."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "An M=1 q8 GEMV with explicit zero points and no bias isolates the zero-point binding layout and epilogue."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "An M=1 q8 GEMV with bias and the schema-default zero point isolates the bias binding layout and epilogue."
       },
       "inputs": {
         "aT": {
       "name": "q8_zero_bias_gemv_m1_tail_n5_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
+        "notes": "With `GEMV_N_COLS = 8` and N=5, q8 unpacking, explicit zero points, and bias run with one live column in the second group."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_zero_bias_naive_fallback_tailK_m3_n6",
       "provenance": {
+        "notes": "K=17 with block size 16 leaves a partial block; M=3 bypasses GEMV and N=6 remains below the tiled floors. The scalar fallback handles q8 unpacking, per-block zero points, bias, and the K tail."
       },
       "inputs": {
         "aT": {
     {
       "name": "q4_prefill_tiled_reg_aligned_m32_k32_n4096_at_workgroup_floor",
       "provenance": {
+        "source": "synthetic",
+        "notes": "Aligned K=32, M=32, and N=4096 produce exactly 64 workgroups, exercising the register-tiled prefill route at its workgroup floor."
       },
       "inputs": {
         "aT": {
     {
       "name": "q4_prefill_tiled_reg_tailk_m32_k33_n512_below_wide_n_floor",
       "provenance": {
+        "notes": "K=33 exercises a padded final block while N=512 places the register-tiled prefill route in its relaxed narrow-output band."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "attrs": { "K": 128, "N": 64, "bits": 8, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 4, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "attrs": { "K": 128, "N": 1024, "bits": 8, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "inputs": {
         "aT": {
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is zeroPointsT; the expected output isolates explicit zero-point decoding from bias handling."
       },
       "attrs": { "K": 128, "N": 64, "bits": 4, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc",
         "test": "independent has_zero_points/has_bias epilogues",
+        "notes": "The only optional input is biasT; zero-point decoding uses the schema default, isolating the bias epilogue."
       },
       "attrs": { "K": 128, "N": 64, "bits": 4, "block_size": 32 }
     },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc",
         "test": "accuracy_level=4 int8 activation quantization (DP4AMatMulQuantize + DP4AMatMulNBits)",
+        "notes": "Accuracy level 4 quantizes A per 128-element block using `max_abs / 127` and signed-normalized rounding, then multiplies q4 weights with packed integer dots. Tolerances account for the permitted activation-quantization error."
       },
       "attrs": { "K": 256, "N": 64, "bits": 4, "block_size": 32, "accuracy_level": 4 }
     },
         "bT": {
           "dtype": "uint8",
           "shape": [16, 2, 32],
+          "data": { "kind": "cycle", "values": [171, 205, 46, 251, 218, 201, 31, 186, 237, 220, 154] }
         },
         "scalesT": { "dtype": "float32", "shape": [16, 2], "data": { "kind": "linspace", "start": 0.02, "end": 0.06 } }
       },
       "provenance": {
         "source": "onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc",
         "test": "accuracy_level=4 int8 activation quantization (DP4AMatMulQuantize + DP4AMatMulNBits)",
+        "notes": "Diverges from the upstream test's inputs (inputs.bT cycle [37, 201, 116, 88, 250, 13, 164, 95, 142, 7, 233] -> cycle [171, 205, 46, 251, 218, 201, 31, 186, 237, 220, 154]); the expected output is recomputed by the CPU reference for the new inputs. Accuracy level 4 quantizes A per 128-element block using `max_abs / 127` and signed-normalized rounding, then multiplies q4 weights with packed integer dots. Tolerances account for the permitted activation-quantization error, so the weight codes sit mostly above the default zero point: a near-zero-mean code cycle cancels the product down to the magnitude of that error, leaving no tolerance that can see a scale mistake."
       },
       "attrs": { "K": 128, "N": 16, "bits": 4, "block_size": 64, "accuracy_level": 4 }
     },
       },
       "outputs": { "yT": { "dtype": "float32", "shape": [3, 6], "tolerance": 0.001 } },
       "attrs": { "K": 17, "N": 6, "bits": 8, "block_size": 16 }
+    },
+    {
+      "name": "q4_prefill_tile4x4_m4_k128_n4100_column_groups_partial_group",
+      "provenance": {
+        "notes": "K=128 at block 32 is sixteen packed words per column, so the small-M kernel runs 16 K-lanes and four column groups (16 columns) per workgroup. N=4100 is 256 full workgroups plus one whose second group starts at column 4096 and whose third and fourth groups lie past N, so the group-valid guard and the partial last group are both exercised at scale; M=4 is one row tile."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [4, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.35 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [4100, 4, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [4100, 4],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [4, 4100], "tolerance": 0.0001 } },
+      "attrs": { "K": 128, "N": 4100, "bits": 4, "block_size": 32 }
+    },
+    {
+      "name": "q8_zero_bias_prefill_tile4x4_m5_k64_n300_column_groups",
+      "provenance": {
+        "notes": "q8 at K=64, block 32: sixteen packed words per column, four 8-bit codes per word, so a word carries four K-adjacent codes of one block. N=300 with sixteen columns per workgroup leaves a final workgroup whose last group holds columns 296..299 and M=5 leaves a second row tile with one live row, so the per-row and per-column guards run together with zero points and bias."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [5, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.27, "scale": 0.4 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [300, 2, 32],
+          "data": { "kind": "cycle", "values": [17, 200, 91, 45, 233, 128, 7, 176, 250, 33, 142, 99, 210, 64, 188] }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [300, 2],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "zeroPointsT": {
+          "dtype": "float32",
+          "shape": [300, 2],
+          "data": { "kind": "cycle", "values": [120.0, 132.0, 128.0, 140.0, 135.0, 118.0, 144.0] }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [300],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.1 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [5, 300], "tolerance": 0.001 } },
+      "attrs": { "K": 64, "N": 300, "bits": 8, "block_size": 32 }
+    },
+    {
+      "name": "q4_gemv_default_zero_m1_k2080_n6_two_vector_trips",
+      "provenance": {
+        "notes": "K=2080 at block size 32 yields 65 sixteen-byte vectors per column, so lane 0 makes a second vector-loop trip and applies the 65th block's scale. N=6 leaves the second column group half empty."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 2080],
+          "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.031 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [6, 65, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [6, 65],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 6], "tolerance": 0.0001 } },
+      "attrs": { "K": 2080, "N": 6, "bits": 4, "block_size": 32 }
+    },
+    {
+      "name": "q8_gemv_zero_bias_m1_k48_block16_word_vectors",
+      "provenance": {
+        "notes": "q8 at block 16 is a sixteen-byte blob, so the GEMV still loads whole sixteen-byte vectors (one per block) while q4 at block 16 would fall to single words; K=48 is three vectors and N=9 leaves a partial column group, with zero points and bias."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 48],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [9, 3, 16],
+          "data": { "kind": "cycle", "values": [17, 200, 91, 45, 233, 128, 7, 176, 250, 33, 142, 99, 210, 64, 188] }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [9, 3],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "zeroPointsT": {
+          "dtype": "float32",
+          "shape": [9, 3],
+          "data": { "kind": "cycle", "values": [120.0, 132.0, 128.0, 140.0, 135.0, 118.0, 144.0] }
+        },
+        "biasT": {
+          "dtype": "float32",
+          "shape": [9],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.1 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 9], "tolerance": 0.0001 } },
+      "attrs": { "K": 48, "N": 9, "bits": 8, "block_size": 16 }
+    },
+    {
+      "name": "q4_gemv_default_zero_m1_block16_single_word_vectors_k40",
+      "provenance": {
+        "notes": "q4 at block 16 is an eight-byte blob, below the sixteen-byte vector, so the GEMV falls to one-word vectors (eight codes each) and K=40 ends half way through the third block: the last vector's trailing codes lie past K and must contribute zero."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 40],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [5, 3, 8],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float32",
+          "shape": [5, 3],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001 } },
+      "attrs": { "K": 40, "N": 5, "bits": 4, "block_size": 16 }
     }
   ]
 }