sync 6fdf6301e2bb

Browse files

Files changed (13) hide show

README.md +10 -10
build/webgpu/bench.json +66 -0
build/webgpu/manifest.json +192 -620
build/webgpu/matmul-nbits-dp4a-quantize.wgsl.jinja +1 -1
build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja +6 -10
build/webgpu/matmul-nbits-q4-dp4a-prefill.wgsl.jinja +3 -4
build/webgpu/matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja +45 -9
build/webgpu/matmul-nbits-q4-prefill-tiled.wgsl.jinja +5 -6
build/webgpu/matmul-nbits-q4-sgmat.wgsl.jinja +2 -2
build/webgpu/matmul-nbits.wgsl.jinja +8 -7
build/webgpu/metadata.json +14 -14
build/webgpu/reduce-axis0-splitk-combine.wgsl.jinja +4 -25
build/webgpu/test.json +185 -15

README.md CHANGED Viewed

@@ -55,14 +55,14 @@ Attributes and default values (overridable per request):
 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
@@ -73,7 +73,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`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-dp4a-quantize.wgsl.jinja`](build/webgpu/matmul-nbits-dp4a-quantize.wgsl.jinja)
 - [`matmul-nbits-gemv-q4.wgsl.jinja`](build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja)
 - [`matmul-nbits-q4-dp4a-prefill.wgsl.jinja`](build/webgpu/matmul-nbits-q4-dp4a-prefill.wgsl.jinja)
@@ -87,7 +87,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## 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.

 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` — Splits the vec4 register-tiled K reduction across dispatch.z, then combines f32 partials and bias. Large 2-bit f16 tiles pin the narrowest subgroup width of a variable range under subgroup-size control; fixed 32-lane paths reuse each packed word, scale, and zero point across four vectors.
+- `prefill_tiled_reg_vec4_default_zero` — Register-blocked prefill with vec4 activation loads. Taller inputs use 128-row tiles. Large 2-bit f16 grids pin the narrowest width of a variable subgroup range under size control, with BK=16 on 16–32 lanes; other tiers use BK=32. Deep 2-bit reductions on fixed 32-lane subgroups unpack four vectors per packed word and reuse its scale and zero point.
+- `prefill_tiled_reg_vec4_splitk_zero_bias` — Splits the vec4 register-tiled K reduction across dispatch.z, then combines f32 partials and bias. Large 2-bit f16 tiles pin the narrowest subgroup width of a variable range under subgroup-size control; fixed 32-lane paths reuse each packed word, scale, and zero point across four vectors.
+- `prefill_tiled_reg_vec4_zero_bias` — Register-blocked prefill with vec4 activation loads. Taller inputs use 128-row tiles. Large 2-bit f16 grids pin the narrowest width of a variable subgroup range under size control, with BK=16 on 16–32 lanes; other tiers use BK=32. Deep 2-bit reductions on fixed 32-lane subgroups unpack four vectors per packed word and reuse its scale and zero point.
+- `prefill_tiled_reg_vec4_splitk_zero_only` — Splits the vec4 register-tiled K reduction across dispatch.z, then combines f32 partials and bias. Large 2-bit f16 tiles pin the narrowest subgroup width of a variable range under subgroup-size control; fixed 32-lane paths reuse each packed word, scale, and zero point across four vectors.
+- `prefill_tiled_reg_vec4_zero_only` — Register-blocked prefill with vec4 activation loads. Taller inputs use 128-row tiles. Large 2-bit f16 grids pin the narrowest width of a variable subgroup range under size control, with BK=16 on 16–32 lanes; other tiers use BK=32. Deep 2-bit reductions on fixed 32-lane subgroups unpack four vectors per packed word and reuse its scale and zero point.
+- `prefill_tiled_reg_vec4_splitk_bias_only` — Splits the vec4 register-tiled K reduction across dispatch.z, then combines f32 partials and bias. Large 2-bit f16 tiles pin the narrowest subgroup width of a variable range under subgroup-size control; fixed 32-lane paths reuse each packed word, scale, and zero point across four vectors.
+- `prefill_tiled_reg_vec4_bias_only` — Register-blocked prefill with vec4 activation loads. Taller inputs use 128-row tiles. Large 2-bit f16 grids pin the narrowest width of a variable subgroup range under size control, with BK=16 on 16–32 lanes; other tiers use BK=32. Deep 2-bit reductions on fixed 32-lane subgroups unpack four vectors per packed word and reuse its scale and zero point.
 ## Device requirements
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`matmul-nbits-dp4a-quantize.wgsl.jinja`](build/webgpu/matmul-nbits-dp4a-quantize.wgsl.jinja)
 - [`matmul-nbits-gemv-q4.wgsl.jinja`](build/webgpu/matmul-nbits-gemv-q4.wgsl.jinja)
 - [`matmul-nbits-q4-dp4a-prefill.wgsl.jinja`](build/webgpu/matmul-nbits-q4-dp4a-prefill.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

build/webgpu/bench.json CHANGED Viewed

@@ -859,6 +859,72 @@
       "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 }
     }
   ]
 }

       "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 }
+    },
+    {
+      "name": "bonsai2-ffn-gate-prefill-m128-q2g128-k5120-n17408-zero1-f16",
+      "preset": "model",
+      "vars": { "M": 128, "K": 5120, "N": 17408, "bits": 2, "blockSize": 128 },
+      "inputs": {
+        "aT": { "shape": [128, 5120], "dtype": "float16", "dist": "normal", "seed": 721, "scale": 0.2 },
+        "bT": { "shape": [17408, 40, 32], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "scalesT": {
+          "shape": [17408, 40],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 722,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        },
+        "zeroPointsT": { "shape": [17408, 40], "dtype": "float16", "dist": "constant", "value": 1 }
+      },
+      "outputs": { "yT": { "shape": [128, 17408], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.M * args.K * args.N" }] },
+      "attrs": { "K": 5120, "N": 17408, "bits": 2, "block_size": 128 }
+    },
+    {
+      "name": "bonsai2-ffn-down-prefill-m128-q2g128-k17408-n5120-zero1-f16",
+      "preset": "model",
+      "vars": { "M": 128, "K": 17408, "N": 5120, "bits": 2, "blockSize": 128 },
+      "inputs": {
+        "aT": { "shape": [128, 17408], "dtype": "float16", "dist": "normal", "seed": 721, "scale": 0.2 },
+        "bT": { "shape": [5120, 136, 32], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "scalesT": {
+          "shape": [5120, 136],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 722,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        },
+        "zeroPointsT": { "shape": [5120, 136], "dtype": "float16", "dist": "constant", "value": 1 }
+      },
+      "outputs": { "yT": { "shape": [128, 5120], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.M * args.K * args.N" }] },
+      "attrs": { "K": 17408, "N": 5120, "bits": 2, "block_size": 128 }
+    },
+    {
+      "name": "bonsai2-attn-qkv-prefill-m128-q2g128-k5120-n10240-zero1-f16",
+      "preset": "model",
+      "vars": { "M": 128, "K": 5120, "N": 10240, "bits": 2, "blockSize": 128 },
+      "inputs": {
+        "aT": { "shape": [128, 5120], "dtype": "float16", "dist": "normal", "seed": 721, "scale": 0.2 },
+        "bT": { "shape": [10240, 40, 32], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "scalesT": {
+          "shape": [10240, 40],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 722,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        },
+        "zeroPointsT": { "shape": [10240, 40], "dtype": "float16", "dist": "constant", "value": 1 }
+      },
+      "outputs": { "yT": { "shape": [128, 10240], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * args.M * args.K * args.N" }] },
+      "attrs": { "K": 5120, "N": 10240, "bits": 2, "block_size": 128 }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -37,10 +37,12 @@
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
     "packedFeature": "device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\")",
     "kBlocksExpected": "ceilDiv(attrs.K, attrs.block_size)",
     "blobSizeExpected": "ceilDiv(attrs.block_size * attrs.bits, 8)",
@@ -73,7 +75,9 @@
     "zeroOnlyEpilogue": "zeroPointsValid and not present.biasT",
     "biasOnlyEpilogue": "not present.zeroPointsT and biasValid",
     "portableWorkgroupFits": "portableWorkgroupSize > 0 and portableWorkgroupSize * 64 <= device.limits.maxComputeWorkgroupStorageSize",
-    "tiledWorkgroupFits": "16 <= device.limits.maxComputeWorkgroupSizeX and 16 <= device.limits.maxComputeWorkgroupSizeY and 256 <= device.limits.maxComputeInvocationsPerWorkgroup and 4096 <= device.limits.maxComputeWorkgroupStorageSize",
     "tiledRegWorkgroupFits": "tiledWorkgroupFits and 16384 <= device.limits.maxComputeWorkgroupStorageSize",
     "mediumTiledRegWorkgroupFits": "tiledWorkgroupFits and 6144 <= device.limits.maxComputeWorkgroupStorageSize",
     "sgmatWorkgroupFits": "sgmatWorkgroupSize <= deviceWorkgroupCap and sgmatWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
@@ -86,8 +90,8 @@
     "mediumTiledRegEligible": "mediumRegisterEligible and mediumTiledRegWorkgroupFits and dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "tiledRegVariantEligible": "largeTiledRegEligible or mediumTiledRegEligible",
     "tiledRegSelectedBK": "tiledRegBK if largeTiledRegEligible else 16",
-    "tiledRegSelectedTileRows": "64 if largeTiledRegEligible else 32",
     "tiledRegSelectedThreadRows": "4 if largeTiledRegEligible else 2",
     "tiledRegSelectedDispatchM": "dispatchM64 if largeTiledRegEligible else dispatchM32",
     "blobWords": "blobSizeExpected / 4",
     "codesPerWord": "32 / attrs.bits",
@@ -100,64 +104,61 @@
     "smallMKLanes": "min(portableWorkgroupSize, pow2ceil(smallMWordsPerCol + 1) / 2)",
     "smallMColGroups": "portableWorkgroupSize / smallMKLanes",
     "smallMDispatchN": "ceilDiv(attrs.N, 4 * smallMColGroups)",
-    "tiledRegVec4TileRows": "128 if aRows >= 256 else 64",
     "tiledRegVec4ThreadRows": "8 if aRows >= 256 else 4",
     "tiledRegVec4DispatchM": "ceilDiv(aRows, tiledRegVec4TileRows)",
     "tiledRegVec4Eligible": "largeTiledRegEligible and attrs.K % 4 == 0 and attrs.block_size % 32 == 0 and tiledRegVec4DispatchM <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "tiledRegSplitTiles": "tiledRegVec4DispatchM * dispatchN64",
     "tiledRegSplitWant": "ceilDiv(tunables.REGISTER_TILE_SPLITK_TARGET_WORKGROUPS, max(1, tiledRegSplitTiles))",
     "tiledRegSplitK": "8 if (tiledRegSplitWant >= 8 and attrs.K >= 4096) else (4 if (tiledRegSplitWant >= 4 and attrs.K >= 2048) else (2 if (tiledRegSplitWant >= 2 and attrs.K >= 1024) else 1))",
     "tiledRegSplitTilesPerSplit": "ceilDiv(ceilDiv(attrs.K, 32), tiledRegSplitK)",
-    "tiledRegSplitEligible": "tiledRegVec4Eligible and tiledRegSplitK >= 2 and tiledRegSplitTiles <= tunables.REGISTER_TILE_SPLITK_MAX_TILES and tiledRegSplitK * aRows * attrs.N * 4 <= device.limits.maxStorageBufferBindingSize and tiledRegSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "B_LEN": "attrs.N * kBlocksExpected * blobSizeExpected / 4",
     "SCALES_LEN": "attrs.N * kBlocksExpected",
     "BIAS_LEN": "attrs.N"
   },
   "bindings": {
-    "a_2": { "arg": "aT", "name": "a", "buffer": "read-only-storage", "elementType": "$aElement" },
-    "b_2": {
-      "arg": "bT",
-      "name": "b",
-      "buffer": "read-only-storage",
-      "elementType": "$bElement",
-      "length": "$B_VEC_LEN"
-    },
-    "scales_2": {
-      "arg": "scalesT",
-      "name": "scales",
-      "buffer": "read-only-storage",
-      "elementType": "$scaleScalar",
-      "length": "$SCALES_LEN"
-    },
-    "y_2": { "arg": "yT", "name": "y", "buffer": "storage", "elementType": "$outputScalar" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "K", "type": "u32", "value": "attrs.K" },
         { "name": "N", "type": "u32", "value": "attrs.N" },
         { "name": "kBlocks", "type": "u32", "value": "dim(shapes.bT, 1)" }
       ]
     },
-    "zero_points": {
-      "arg": "zeroPointsT",
-      "buffer": "read-only-storage",
-      "elementType": "$aScalar",
-      "length": "$SCALES_LEN"
-    },
-    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$BIAS_LEN" },
-    "a_3": { "arg": "aT", "name": "a", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "b_3": { "arg": "bT", "name": "b", "buffer": "read-only-storage", "elementType": "$bScalar", "length": "$B_LEN" },
-    "a_4": { "arg": "aT", "name": "a", "buffer": "read-only-storage", "elementType": "$aVec4Element" },
-    "y_3": { "scratch": "partials", "name": "y", "buffer": "storage", "elementType": "f32" },
     "partials": { "buffer": "read-only-storage", "elementType": "f32" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
-      "struct": [{ "name": "cols", "type": "u32", "value": "aRows * attrs.N" }]
-    },
-    "params_3": {
-      "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "rows", "type": "u32", "value": "aRows" },
         { "name": "K", "type": "u32", "value": "attrs.K" },
@@ -172,18 +173,9 @@
     {
       "id": "q4_dp4a_prefill",
       "priority": 19,
-      "when": ["packedFeature", "commonShapeValid", "defaultEpilogue", "attrs.bits == 4", "attrs.accuracy_level == 4", "tensorDtypes.aT == \"float32\"", "attrs.block_size % 32 == 0", "attrs.K % 128 == 0", "attrs.N % 16 == 0", "aRows >= 32", "ceilDiv(attrs.N, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(aRows, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "16 <= device.limits.maxComputeWorkgroupSizeX", "16 <= device.limits.maxComputeWorkgroupSizeY", "64 <= device.limits.maxComputeWorkgroupSizeX", "256 <= device.limits.maxComputeInvocationsPerWorkgroup", "4608 <= device.limits.maxComputeWorkgroupStorageSize"],
       "demoteWhen": ["has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"maxwell\"", "device.features.has(\"chromium-experimental-subgroup-matrix\")", "device.adapterInfo.vendor == \"apple\""],
-      "derive": {
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "vec4Count": "aRows * attrs.K / 4",
-        "blockCount": "aRows * attrs.K / 128"
-      },
       "intermediates": [
         { "id": "aQuant", "dtype": "uint32", "shape": "[aRows * attrs.K / 4]" },
         { "id": "aScales", "dtype": "float32", "shape": "[aRows * attrs.K / 128]" }
@@ -199,8 +191,8 @@
             { "scratch": "aScales", "name": "a_scales", "elementType": "f32" }
           ],
           "dispatch": {
-            "x": "min(ceilDiv((aRows * attrs.K / 4), (64)), 65535)",
-            "y": "ceilDiv(ceilDiv((aRows * attrs.K / 4), (64)), 65535)",
             "z": 1
           }
         },
@@ -226,33 +218,22 @@
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
-        "hasZero": false,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
-        "actVec4": "gemvActVec4",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
-          "bindings": ["a_2", "b_2", "scales_2", "y_2", "params"],
-          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -266,21 +247,6 @@
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
@@ -292,7 +258,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
@@ -301,28 +267,18 @@
       "id": "prefill_tiled_reg_vec4_splitk_default_zero",
       "priority": 17,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
@@ -336,22 +292,17 @@
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "y_3"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": {
-            "op": "\"sum\"",
-            "outputF16": "tensorDtypes.aT == \"float16\"",
-            "intMode": false,
-            "addBias": false
-          },
-          "bindings": ["partials", "y_2", "params_2"],
           "dispatch": {
-            "x": "min(ceilDiv((aRows * attrs.N), (256)), 65535)",
-            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (256)), 65535)",
             "z": 1
           }
         }
@@ -362,27 +313,16 @@
       "priority": 16,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
-        "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
@@ -390,8 +330,8 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "y_2"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM" }
         }
       ]
     },
@@ -400,21 +340,6 @@
       "priority": 15,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
@@ -426,7 +351,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
@@ -435,28 +360,12 @@
       "id": "prefill_tiled_default_zero",
       "priority": 14,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
-      "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "y_2"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
@@ -466,27 +375,16 @@
       "priority": 13,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
-        "colGroups": "smallMColGroups",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "y_2", "params_3"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
@@ -499,23 +397,12 @@
       "id": "default_zero",
       "priority": 0,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "portableWorkgroupFits"],
-      "derive": {
-        "hasZero": false,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "y_2", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
@@ -531,33 +418,22 @@
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
-        "hasZero": true,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
-        "actVec4": "gemvActVec4",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
-          "bindings": ["a_2", "b_2", "scales_2", "zero_points", "bias", "y_2", "params"],
-          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -571,21 +447,6 @@
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
@@ -597,7 +458,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "bias", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
@@ -606,28 +467,18 @@
       "id": "prefill_tiled_reg_vec4_splitk_zero_bias",
       "priority": 17,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
@@ -641,22 +492,17 @@
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "zero_points", "y_3"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": {
-            "op": "\"sum\"",
-            "outputF16": "tensorDtypes.aT == \"float16\"",
-            "intMode": false,
-            "addBias": true
-          },
-          "bindings": ["partials", "bias", "y_2", "params_2"],
           "dispatch": {
-            "x": "min(ceilDiv((aRows * attrs.N), (256)), 65535)",
-            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (256)), 65535)",
             "z": 1
           }
         }
@@ -667,27 +513,16 @@
       "priority": 16,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
-        "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
@@ -695,8 +530,8 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "zero_points", "bias", "y_2"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM" }
         }
       ]
     },
@@ -705,21 +540,6 @@
       "priority": 15,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
@@ -731,7 +551,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "bias", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
@@ -740,28 +560,12 @@
       "id": "prefill_tiled_zero_bias",
       "priority": 14,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
-      "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "bias", "y_2"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
@@ -771,27 +575,16 @@
       "priority": 13,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
-        "colGroups": "smallMColGroups",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "bias", "y_2", "params_3"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
@@ -804,23 +597,12 @@
       "id": "zero_bias",
       "priority": 0,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "portableWorkgroupFits"],
-      "derive": {
-        "hasZero": true,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "bias", "y_2", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
@@ -836,33 +618,22 @@
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
-        "hasZero": true,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
-        "actVec4": "gemvActVec4",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
-          "bindings": ["a_2", "b_2", "scales_2", "zero_points", "y_2", "params"],
-          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -876,21 +647,6 @@
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
@@ -902,7 +658,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
@@ -911,28 +667,18 @@
       "id": "prefill_tiled_reg_vec4_splitk_zero_only",
       "priority": 17,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
@@ -946,22 +692,17 @@
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "zero_points", "y_3"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": {
-            "op": "\"sum\"",
-            "outputF16": "tensorDtypes.aT == \"float16\"",
-            "intMode": false,
-            "addBias": false
-          },
-          "bindings": ["partials", "y_2", "params_2"],
           "dispatch": {
-            "x": "min(ceilDiv((aRows * attrs.N), (256)), 65535)",
-            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (256)), 65535)",
             "z": 1
           }
         }
@@ -972,27 +713,16 @@
       "priority": 16,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
-        "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
@@ -1000,8 +730,8 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "zero_points", "y_2"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM" }
         }
       ]
     },
@@ -1010,21 +740,6 @@
       "priority": 15,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
@@ -1036,7 +751,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
@@ -1045,28 +760,12 @@
       "id": "prefill_tiled_zero_only",
       "priority": 14,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
-      "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "y_2"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
@@ -1076,27 +775,16 @@
       "priority": 13,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
-        "colGroups": "smallMColGroups",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "y_2", "params_3"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
@@ -1109,23 +797,12 @@
       "id": "zero_only",
       "priority": 0,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"],
-      "derive": {
-        "hasZero": true,
-        "hasBias": false,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "zero_points", "y_2", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
@@ -1141,33 +818,22 @@
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
-        "hasZero": false,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
-        "actVec4": "gemvActVec4",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
-          "bindings": ["a_2", "b_2", "scales_2", "bias", "y_2", "params"],
-          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
@@ -1181,21 +847,6 @@
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
@@ -1207,7 +858,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "bias", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
@@ -1216,28 +867,18 @@
       "id": "prefill_tiled_reg_vec4_splitk_bias_only",
       "priority": 17,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
@@ -1251,22 +892,17 @@
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "y_3"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
-          "derive": {
-            "op": "\"sum\"",
-            "outputF16": "tensorDtypes.aT == \"float16\"",
-            "intMode": false,
-            "addBias": true
-          },
-          "bindings": ["partials", "bias", "y_2", "params_2"],
           "dispatch": {
-            "x": "min(ceilDiv((aRows * attrs.N), (256)), 65535)",
-            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (256)), 65535)",
             "z": 1
           }
         }
@@ -1277,27 +913,16 @@
       "priority": 16,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
-        "bk": 32,
-        "tileRows": "tiledRegVec4TileRows",
-        "tileCols": 64,
-        "threadRows": "tiledRegVec4ThreadRows",
-        "threadCols": 4,
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
@@ -1305,8 +930,8 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_4", "b_3", "scales_2", "bias", "y_2"],
-          "dispatch": { "x": "dispatchN64", "y": "tiledRegVec4DispatchM" }
         }
       ]
     },
@@ -1315,21 +940,6 @@
       "priority": 15,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\"",
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
@@ -1341,7 +951,7 @@
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "bias", "y_2"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
@@ -1350,28 +960,12 @@
       "id": "prefill_tiled_bias_only",
       "priority": 14,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
-      "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "M": "aRows",
-        "K": "attrs.K",
-        "N": "attrs.N",
-        "kBlocks": "dim(shapes.bT, 1)",
-        "blockSize": "attrs.block_size",
-        "blobSize": "dim(shapes.bT, 2)",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "bias", "y_2"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
@@ -1381,27 +975,16 @@
       "priority": 13,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
-        "colGroups": "smallMColGroups",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "bias", "y_2", "params_3"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
@@ -1414,23 +997,12 @@
       "id": "bias_only",
       "priority": 0,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"],
-      "derive": {
-        "hasZero": false,
-        "hasBias": true,
-        "workgroupSize": "portableWorkgroupSize",
-        "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bScalar": "\"u32\"",
-        "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-        "bits": "attrs.bits",
-        "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-        "usesF16": "tensorDtypes.aT == \"float16\""
-      },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
-          "bindings": ["a_3", "b_3", "scales_2", "bias", "y_2", "params_3"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",

   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
+    "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
     "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
     "wave32Effective": "wave32Adapter or pinSubgroupSize32",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
     "packedFeature": "device.wgslLanguageFeatures.has(\"packed_4x8_integer_dot_product\")",
     "kBlocksExpected": "ceilDiv(attrs.K, attrs.block_size)",
     "blobSizeExpected": "ceilDiv(attrs.block_size * attrs.bits, 8)",
     "zeroOnlyEpilogue": "zeroPointsValid and not present.biasT",
     "biasOnlyEpilogue": "not present.zeroPointsT and biasValid",
     "portableWorkgroupFits": "portableWorkgroupSize > 0 and portableWorkgroupSize * 64 <= device.limits.maxComputeWorkgroupStorageSize",
+    "tiledWorkgroupSide": 16,
+    "dp4aQuantizeWorkgroupSize": 64,
+    "tiledWorkgroupFits": "tiledWorkgroupSide <= device.limits.maxComputeWorkgroupSizeX and tiledWorkgroupSide <= device.limits.maxComputeWorkgroupSizeY and tiledWorkgroupSide * tiledWorkgroupSide <= device.limits.maxComputeInvocationsPerWorkgroup and 4096 <= device.limits.maxComputeWorkgroupStorageSize",
     "tiledRegWorkgroupFits": "tiledWorkgroupFits and 16384 <= device.limits.maxComputeWorkgroupStorageSize",
     "mediumTiledRegWorkgroupFits": "tiledWorkgroupFits and 6144 <= device.limits.maxComputeWorkgroupStorageSize",
     "sgmatWorkgroupFits": "sgmatWorkgroupSize <= deviceWorkgroupCap and sgmatWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
     "mediumTiledRegEligible": "mediumRegisterEligible and mediumTiledRegWorkgroupFits and dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "tiledRegVariantEligible": "largeTiledRegEligible or mediumTiledRegEligible",
     "tiledRegSelectedBK": "tiledRegBK if largeTiledRegEligible else 16",
     "tiledRegSelectedThreadRows": "4 if largeTiledRegEligible else 2",
+    "tiledRegSelectedTileRows": "tiledRegSelectedThreadRows * tiledWorkgroupSide",
     "tiledRegSelectedDispatchM": "dispatchM64 if largeTiledRegEligible else dispatchM32",
     "blobWords": "blobSizeExpected / 4",
     "codesPerWord": "32 / attrs.bits",
     "smallMKLanes": "min(portableWorkgroupSize, pow2ceil(smallMWordsPerCol + 1) / 2)",
     "smallMColGroups": "portableWorkgroupSize / smallMKLanes",
     "smallMDispatchN": "ceilDiv(attrs.N, 4 * smallMColGroups)",
     "tiledRegVec4ThreadRows": "8 if aRows >= 256 else 4",
+    "tiledRegVec4TileRows": "tiledRegVec4ThreadRows * tiledWorkgroupSide",
     "tiledRegVec4DispatchM": "ceilDiv(aRows, tiledRegVec4TileRows)",
     "tiledRegVec4Eligible": "largeTiledRegEligible and attrs.K % 4 == 0 and attrs.block_size % 32 == 0 and tiledRegVec4DispatchM <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "tiledRegSubgroupPin": "device.adapterInfo.subgroupMinSize if (device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMinSize <= tiledWorkgroupSide and attrs.bits == 2 and tensorDtypes.aT == \"float16\" and tiledRegVec4Eligible and aRows >= 2 * tiledRegVec4TileRows) else 0",
+    "q2WideMicroPreferred": "variableSubgroup16To32 and device.features.has(\"subgroup-size-control\") and attrs.bits == 2 and tensorDtypes.aT == \"float16\" and tiledRegVec4ThreadRows == 4 and aRows >= 2 * tiledRegVec4TileRows and device.limits.maxComputeWorkgroupSizeX >= 32 and device.limits.maxComputeWorkgroupSizeY >= 8 and device.limits.maxComputeInvocationsPerWorkgroup >= 256",
     "tiledRegSplitTiles": "tiledRegVec4DispatchM * dispatchN64",
+    "q2NarrowBkPreferred": "variableSubgroup16To32 and attrs.bits == 2 and tensorDtypes.aT == \"float16\" and attrs.block_size >= 128 and attrs.block_size % 128 == 0 and attrs.K % attrs.block_size == 0 and tiledRegVec4TileRows == 64 and aRows >= 2 * tiledRegVec4TileRows and tiledRegSplitTiles >= tunables.REGISTER_TILE_SPLITK_MAX_TILES / 2",
     "tiledRegSplitWant": "ceilDiv(tunables.REGISTER_TILE_SPLITK_TARGET_WORKGROUPS, max(1, tiledRegSplitTiles))",
     "tiledRegSplitK": "8 if (tiledRegSplitWant >= 8 and attrs.K >= 4096) else (4 if (tiledRegSplitWant >= 4 and attrs.K >= 2048) else (2 if (tiledRegSplitWant >= 2 and attrs.K >= 1024) else 1))",
     "tiledRegSplitTilesPerSplit": "ceilDiv(ceilDiv(attrs.K, 32), tiledRegSplitK)",
+    "tiledRegSplitEligible": "tiledRegVec4Eligible and tiledRegSplitK >= 2 and tiledRegSplitTiles <= tunables.REGISTER_TILE_SPLITK_MAX_TILES and not (q2NarrowBkPreferred and attrs.block_size >= 256) and tiledRegSplitK * aRows * attrs.N * 4 <= device.limits.maxStorageBufferBindingSize and tiledRegSplitK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "hasZero": "present.zeroPointsT",
+    "hasBias": "present.biasT",
+    "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "bScalar": "\"u32\"",
+    "scaleScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "outputScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "usesF16": "tensorDtypes.aT == \"float16\"",
+    "M": "aRows",
+    "K": "attrs.K",
+    "N": "attrs.N",
+    "kBlocks": "dim(shapes.bT, 1)",
+    "blockSize": "attrs.block_size",
+    "blobSize": "dim(shapes.bT, 2)",
+    "bits": "attrs.bits",
+    "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
+    "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
+    "workgroupSize": "portableWorkgroupSize",
     "B_LEN": "attrs.N * kBlocksExpected * blobSizeExpected / 4",
     "SCALES_LEN": "attrs.N * kBlocksExpected",
     "BIAS_LEN": "attrs.N"
   },
   "bindings": {
+    "a_a_t": { "arg": "aT", "name": "a", "elementType": "$aElement" },
+    "b_b_t": { "arg": "bT", "name": "b", "elementType": "$bElement", "length": "$B_VEC_LEN" },
+    "scales_main": { "arg": "scalesT", "name": "scales", "elementType": "$scaleScalar", "length": "$SCALES_LEN" },
+    "y_y_t": { "arg": "yT", "name": "y", "elementType": "$outputScalar" },
     "params": {
       "struct": [
         { "name": "K", "type": "u32", "value": "attrs.K" },
         { "name": "N", "type": "u32", "value": "attrs.N" },
         { "name": "kBlocks", "type": "u32", "value": "dim(shapes.bT, 1)" }
       ]
     },
+    "zero_points": { "arg": "zeroPointsT", "elementType": "$aScalar", "length": "$SCALES_LEN" },
+    "bias": { "arg": "biasT", "elementType": "$aScalar", "length": "$BIAS_LEN" },
+    "a_main": { "arg": "aT", "name": "a", "elementType": "$aScalar" },
+    "b_main": { "arg": "bT", "name": "b", "elementType": "$bScalar", "length": "$B_LEN" },
+    "a_partial": { "arg": "aT", "name": "a", "elementType": "$aVec4Element" },
+    "y_f32": { "scratch": "partials", "name": "y", "elementType": "f32" },
     "partials": { "buffer": "read-only-storage", "elementType": "f32" },
+    "params_cols": { "name": "params", "struct": [{ "name": "cols", "type": "u32", "value": "aRows * attrs.N" }] },
+    "params_main": {
       "name": "params",
       "struct": [
         { "name": "rows", "type": "u32", "value": "aRows" },
         { "name": "K", "type": "u32", "value": "attrs.K" },
     {
       "id": "q4_dp4a_prefill",
       "priority": 19,
+      "when": ["packedFeature", "commonShapeValid", "defaultEpilogue", "attrs.bits == 4", "attrs.accuracy_level == 4", "tensorDtypes.aT == \"float32\"", "attrs.block_size % 32 == 0", "attrs.K % 128 == 0", "attrs.N % 16 == 0", "aRows >= 32", "ceilDiv(attrs.N, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "ceilDiv(aRows, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits", "dp4aQuantizeWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "4608 <= device.limits.maxComputeWorkgroupStorageSize"],
       "demoteWhen": ["has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"maxwell\"", "device.features.has(\"chromium-experimental-subgroup-matrix\")", "device.adapterInfo.vendor == \"apple\""],
+      "derive": { "vec4Count": "aRows * attrs.K / 4", "blockCount": "aRows * attrs.K / 128" },
       "intermediates": [
         { "id": "aQuant", "dtype": "uint32", "shape": "[aRows * attrs.K / 4]" },
         { "id": "aScales", "dtype": "float32", "shape": "[aRows * attrs.K / 128]" }
             { "scratch": "aScales", "name": "a_scales", "elementType": "f32" }
           ],
           "dispatch": {
+            "x": "min(ceilDiv((aRows * attrs.K / 4), (dp4aQuantizeWorkgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((aRows * attrs.K / 4), (dp4aQuantizeWorkgroupSize)), 65535)",
             "z": 1
           }
         },
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
+        "actVec4": "gemvActVec4"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
+          "bindings": ["a_a_t", "b_b_t", "scales_main", "y_y_t", "params"],
+          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 }
         }
       ]
     },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
       "id": "prefill_tiled_reg_vec4_splitk_default_zero",
       "priority": 17,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
+      "demoteWhen": ["variableSubgroup16To32 and attrs.bits >= 4 and tiledRegSplitK >= 4 and aRows >= 2 * tiledRegVec4TileRows"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
         "bk": 32,
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "y_f32"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "tensorDtypes.aT == \"float16\"", "addBias": "hasBias" },
+          "bindings": ["partials", "y_y_t", "params_cols"],
           "dispatch": {
+            "x": "min(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
             "z": 1
           }
         }
       "priority": 16,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
+        "bk": "16 if q2NarrowBkPreferred else 32",
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "y_y_t"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)" }
         }
       ]
     },
       "priority": 15,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
       "id": "prefill_tiled_default_zero",
       "priority": 14,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "y_y_t"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
       "priority": 13,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
+        "colGroups": "smallMColGroups"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
       "id": "default_zero",
       "priority": 0,
       "when": ["commonShapeValid", "defaultEpilogue", "bitsSupported", "portableWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
+        "actVec4": "gemvActVec4"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
+          "bindings": ["a_a_t", "b_b_t", "scales_main", "zero_points", "bias", "y_y_t", "params"],
+          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 }
         }
       ]
     },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
       "id": "prefill_tiled_reg_vec4_splitk_zero_bias",
       "priority": 17,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
+      "demoteWhen": ["variableSubgroup16To32 and attrs.bits >= 4 and tiledRegSplitK >= 4 and aRows >= 2 * tiledRegVec4TileRows"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
         "bk": 32,
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "zero_points", "y_f32"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "tensorDtypes.aT == \"float16\"", "addBias": "hasBias" },
+          "bindings": ["partials", "bias", "y_y_t", "params_cols"],
           "dispatch": {
+            "x": "min(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
             "z": 1
           }
         }
       "priority": 16,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
+        "bk": "16 if q2NarrowBkPreferred else 32",
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "zero_points", "bias", "y_y_t"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)" }
         }
       ]
     },
       "priority": 15,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
       "id": "prefill_tiled_zero_bias",
       "priority": 14,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
       "priority": 13,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
+        "colGroups": "smallMColGroups"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "bias", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
       "id": "zero_bias",
       "priority": 0,
       "when": ["commonShapeValid", "zeroBiasEpilogue", "bitsSupported", "portableWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "bias", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
+        "actVec4": "gemvActVec4"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
+          "bindings": ["a_a_t", "b_b_t", "scales_main", "zero_points", "y_y_t", "params"],
+          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 }
         }
       ]
     },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
       "id": "prefill_tiled_reg_vec4_splitk_zero_only",
       "priority": 17,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
+      "demoteWhen": ["variableSubgroup16To32 and attrs.bits >= 4 and tiledRegSplitK >= 4 and aRows >= 2 * tiledRegVec4TileRows"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
         "bk": 32,
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "zero_points", "y_f32"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "tensorDtypes.aT == \"float16\"", "addBias": "hasBias" },
+          "bindings": ["partials", "y_y_t", "params_cols"],
           "dispatch": {
+            "x": "min(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
             "z": 1
           }
         }
       "priority": 16,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
+        "bk": "16 if q2NarrowBkPreferred else 32",
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "zero_points", "y_y_t"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)" }
         }
       ]
     },
       "priority": 15,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
       "id": "prefill_tiled_zero_only",
       "priority": 14,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "y_y_t"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
       "priority": 13,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
+        "colGroups": "smallMColGroups"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
       "id": "zero_only",
       "priority": 0,
       "when": ["commonShapeValid", "zeroOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "zero_points", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
       "derive": {
         "gemvNCols": "tunables.GEMV_N_COLS",
         "useSubgroups": "device.features.has(\"subgroups\")",
         "aElement": "(\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\") if gemvActVec4 else (\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\")",
         "bElement": "\"vec4<u32>\" if gemvVecWords == 4 else \"u32\"",
         "B_VEC_LEN": "attrs.N * gemvVecPerCol",
         "vecWords": "gemvVecWords",
         "codesPerVec": "gemvCodesPerVec",
         "codesPerVec4": "gemvCodesPerVec / 4",
         "vecPerBlock": "gemvVecPerBlock",
         "vecPerCol": "gemvVecPerCol",
+        "actVec4": "gemvActVec4"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-gemv-q4.wgsl.jinja",
+          "bindings": ["a_a_t", "b_b_t", "scales_main", "bias", "y_y_t", "params"],
+          "dispatch": { "x": "min(gemvDispatchN, 65535)", "y": "ceilDiv(gemvDispatchN, 65535)", "z": 1 }
         }
       ]
     },
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "derive": {
         "tileRows": "sgmatTileRows",
         "workgroupSize": "sgmatWorkgroupSize",
         "rowSubtiles": "sgmatRowSubtiles",
         {
           "id": "main",
           "shader": "matmul-nbits-q4-sgmat.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "sgmatDispatchM" }
         }
       ]
       "id": "prefill_tiled_reg_vec4_splitk_bias_only",
       "priority": 17,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegSplitEligible"],
+      "demoteWhen": ["variableSubgroup16To32 and attrs.bits >= 4 and tiledRegSplitK >= 4 and aRows >= 2 * tiledRegVec4TileRows"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
         "bk": 32,
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true,
         "splitK": "tiledRegSplitK",
         {
           "id": "partial",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "y_f32"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)", "z": "tiledRegSplitK" }
         },
         {
           "id": "combine",
           "shader": "reduce-axis0-splitk-combine.wgsl.jinja",
+          "derive": { "op": "\"sum\"", "outputF16": "tensorDtypes.aT == \"float16\"", "addBias": "hasBias" },
+          "bindings": ["partials", "bias", "y_y_t", "params_cols"],
           "dispatch": {
+            "x": "min(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((aRows * attrs.N), (workgroupSize)), 65535)",
             "z": 1
           }
         }
       "priority": 16,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVec4Eligible"],
       "derive": {
         "aVec4Element": "\"vec4<f16>\" if tensorDtypes.aT == \"float16\" else \"vec4<f32>\"",
+        "bk": "16 if q2NarrowBkPreferred else 32",
+        "q2ChunkLoads": "attrs.bits == 2 and attrs.K % 16 == 0 and attrs.block_size % 16 == 0 and subgroupsWave32 and kBlocksExpected >= device.adapterInfo.subgroupMinSize and dispatchN64 >= device.adapterInfo.subgroupMinSize and aRows >= 2 * tiledRegVec4TileRows",
+        "tileWorkgroupX": "32 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "tileWorkgroupY": "8 if q2WideMicroPreferred else tiledWorkgroupSide",
+        "threadRows": "8 if q2WideMicroPreferred else tiledRegVec4ThreadRows",
+        "threadCols": "2 if q2WideMicroPreferred else 4",
+        "tileRows": "threadRows * tileWorkgroupY",
+        "tileCols": "threadCols * tileWorkgroupX",
+        "tileSubgroupPin": "32 if q2WideMicroPreferred else tiledRegSubgroupPin",
         "alignedBlockLoads": true,
         "aVec4Loads": true
       },
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_partial", "b_main", "scales_main", "bias", "y_y_t"],
+          "dispatch": { "x": "dispatchN64", "y": "ceilDiv(aRows, tileRows)" }
         }
       ]
     },
       "priority": 15,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "dispatchN64 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledRegVariantEligible"],
       "derive": {
         "bk": "tiledRegSelectedBK",
         "tileRows": "tiledRegSelectedTileRows",
         "tileCols": 64,
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN64", "y": "tiledRegSelectedDispatchM" }
         }
       ]
       "id": "prefill_tiled_bias_only",
       "priority": 14,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 64", "dispatchN32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dispatchM32 <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "tiledWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tiled.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "bias", "y_y_t"],
           "dispatch": { "x": "dispatchN32", "y": "dispatchM32" }
         }
       ]
       "priority": 13,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "attrs.K % attrs.block_size == 0", "aRows >= 2", "portableWorkgroupFits"],
       "derive": {
         "wordsPerCol": "smallMWordsPerCol",
         "wordsPerBlock": "blobWords",
         "kLanes": "smallMKLanes",
+        "colGroups": "smallMColGroups"
       },
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "bias", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(smallMDispatchN, DISPATCH_FOLD_WIDTH)",
             "y": "min(ceilDiv(aRows, 4), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
       "id": "bias_only",
       "priority": 0,
       "when": ["commonShapeValid", "biasOnlyEpilogue", "bitsSupported", "portableWorkgroupFits"],
+      "derive": {},
       "passes": [
         {
           "id": "main",
           "shader": "matmul-nbits.wgsl.jinja",
+          "bindings": ["a_main", "b_main", "scales_main", "bias", "y_y_t", "params_main"],
           "dispatch": {
             "x": "min(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numel(shapes.yT)), (workgroupSize)), 65535)",

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

@@ -6,7 +6,7 @@
 // multiple of 128, so blocks never straddle rows and the flat layout is exact.
 const VEC4_COUNT: u32 = {{ vec4Count }}u;
 const BLOCK_COUNT: u32 = {{ blockCount }}u;
-const WG: u32 = 64u;
 var<workgroup> maxAbs: array<f32, WG>;

 // multiple of 128, so blocks never straddle rows and the flat layout is exact.
 const VEC4_COUNT: u32 = {{ vec4Count }}u;
 const BLOCK_COUNT: u32 = {{ blockCount }}u;
+const WG: u32 = {{ dp4aQuantizeWorkgroupSize }}u;
 var<workgroup> maxAbs: array<f32, WG>;

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

@@ -46,15 +46,6 @@ fn zero_point({% if hasZero %}n: u32, block: u32{% endif %}) -> f32 {
 {% 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>,
@@ -108,7 +99,12 @@ fn main(
       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 %}

 {% endif %}
 }
 @compute @workgroup_size({{ workgroupSize }}, 1, 1)
 fn main(
   @builtin(workgroup_id) wid: vec3<u32>,
       let scale = f32(scales[n * params.kBlocks + block]);
       let zero = zero_point({% if hasZero %}n, block{% endif %});
       var dot = 0.0;
+{% 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 %}
       acc{{ sfx }}.{{ comps[c] }} = acc{{ sfx }}.{{ comps[c] }} + (dot - zero * asum) * scale;
     }
 {% endfor %}

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

@@ -5,7 +5,6 @@ fn dot4_packed(a_word: u32, b_word: u32) -> i32 {
   return dot4I8Packed(a_word, b_word);
 }
 // com.microsoft.MatMulNBits q4 prefill with int8-quantized activations
 // (accuracy_level 4). A arrives pre-quantized as packed int8 words with one
 // scale per 128-element block; each weight nibble is rebiased by -8 and packed
@@ -27,14 +26,14 @@ var<workgroup> tB: array<array<u32, 8u>, 64u>;
 var<workgroup> tAscale: array<f32, 64u>;
 var<workgroup> tBscale: array<f32, 64u>;
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * 64u;
   let nBase = wg.x * 64u;
-  let li = lid.y * 16u + lid.x;
   var acc: array<f32, 16u>;
   for (var t = 0u; t < 16u; t = t + 1u) { acc[t] = 0.0; }
@@ -46,7 +45,7 @@ fn main(
     // Stage 64 rows x 8 packed A words and 64 cols x 8 packed B words; each of
     // the 256 threads loads two of each.
     for (var e = 0u; e < 2u; e = e + 1u) {
-      let idx = li + e * 256u;
       let r = idx / 8u;
       let w = idx % 8u;
       let am = mBase + r;

   return dot4I8Packed(a_word, b_word);
 }
 // com.microsoft.MatMulNBits q4 prefill with int8-quantized activations
 // (accuracy_level 4). A arrives pre-quantized as packed int8 words with one
 // scale per 128-element block; each weight nibble is rebiased by -8 and packed
 var<workgroup> tAscale: array<f32, 64u>;
 var<workgroup> tBscale: array<f32, 64u>;
+@compute @workgroup_size({{ tiledWorkgroupSide }}, {{ tiledWorkgroupSide }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * 64u;
   let nBase = wg.x * 64u;
+  let li = lid.y * {{ tiledWorkgroupSide }}u + lid.x;
   var acc: array<f32, 16u>;
   for (var t = 0u; t < 16u; t = t + 1u) { acc[t] = 0.0; }
     // Stage 64 rows x 8 packed A words and 64 cols x 8 packed B words; each of
     // the 256 threads loads two of each.
     for (var e = 0u; e < 2u; e = e + 1u) {
+      let idx = li + e * {{ tiledWorkgroupSide * tiledWorkgroupSide }}u;
       let r = idx / 8u;
       let w = idx % 8u;
       let am = mBase + r;

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

@@ -1,3 +1,7 @@
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier — register-blocked.
@@ -6,8 +10,9 @@
 // blocking. Both tiles are indexed by their own output axis and group four K
 // values per vector word, so the micro-tile accumulates through dot() and the
 // column-owning loader writes whole words instead of a BN-strided column.
-// The full geometry computes a 4x4 micro-tile over a 64x64 output tile. The
-// portable geometry computes 2x4 over 32x64 to bound the per-lane accumulator
 // set. K_TILE specializes the K tile. For standard 32/64-element quant blocks,
 // one lane owns one output column and the full BK slice: scale and zero are
 // loaded once, and each stored byte is read once for the K-adjacent codes it
@@ -24,8 +29,10 @@ const BM: u32 = {{ tileRows }}u;
 const BN: u32 = {{ tileCols }}u;
 const TM: u32 = {{ threadRows }}u;
 const TN: u32 = {{ threadCols }}u;
-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 %}
@@ -61,14 +68,12 @@ fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
   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, kBlocks="KBLOCKS", blobSize="BLOB_SIZE") }}
 {% endif %}
 {% macro zero_of(blockExpr) %}{% if hasZero %}f32(zero_points[bn * KBLOCKS + {{ blockExpr }}]){% else %}{{ defaultZero }}{% endif %}{% endmacro %}
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
@@ -115,6 +120,36 @@ fn main(
       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
@@ -162,6 +197,7 @@ fn main(
       }
       tileB[bc][kv] = word;
     }
 {% else %}
     // Quant blocks that do not contain a whole BK tile use the element-wise loader.
     for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + WG_THREADS) {

+{% set subgroupPin = tileSubgroupPin if tileSubgroupPin is defined else 0 %}
+{% if subgroupPin %}
+enable subgroup_size_control;
+{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.MatMulNBits q4/q8 prefill, no-subgroup-matrix tier — register-blocked.
 // blocking. Both tiles are indexed by their own output axis and group four K
 // values per vector word, so the micro-tile accumulates through dot() and the
 // column-owning loader writes whole words instead of a BN-strided column.
+// The full geometry computes a 4x4 micro-tile over a 64x64 output tile. A
+// device-selected 32x8 workgroup can use an 8x2 micro-tile on that same output
+// tile. The portable geometry computes 2x4 over 32x64 to bound the accumulator
 // set. K_TILE specializes the K tile. For standard 32/64-element quant blocks,
 // one lane owns one output column and the full BK slice: scale and zero are
 // loaded once, and each stored byte is read once for the K-adjacent codes it
 const BN: u32 = {{ tileCols }}u;
 const TM: u32 = {{ threadRows }}u;
 const TN: u32 = {{ threadCols }}u;
+{% set wgX = tileWorkgroupX if tileWorkgroupX is defined else tiledWorkgroupSide %}
+{% set wgY = tileWorkgroupY if tileWorkgroupY is defined else tiledWorkgroupSide %}
+const WG_X: u32 = {{ wgX }}u;
+const WG_Y: u32 = {{ wgY }}u;
 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 %}
   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, kBlocks="KBLOCKS", blobSize="BLOB_SIZE") }}
 {% endif %}
 {% macro zero_of(blockExpr) %}{% if hasZero %}f32(zero_points[bn * KBLOCKS + {{ blockExpr }}]){% else %}{{ defaultZero }}{% endif %}{% endmacro %}
+@compute @workgroup_size({{ wgX }}, {{ wgY }}, 1){{ (" @subgroup_size(" ~ subgroupPin ~ ")") if subgroupPin else "" }}
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
       tileA[ar][ac4] = aWord;
     }
 {% if alignedBlockLoads %}
+{% if q2ChunkLoads is defined and q2ChunkLoads and bits == 2 and bk % 16 == 0 %}
+    // A 32-bit word carries sixteen adjacent 2-bit codes. Decode it once per
+    // column and quant chunk, reusing the scale and zero for four staged vec4s.
+    for (var idx: u32 = li; idx < BN * (BK / 16u); idx = idx + WG_THREADS) {
+      let chunksPerCol = BK / 16u;
+      let bc = idx / chunksPerCol;
+      let chunk = idx % chunksPerCol;
+      let bn = nBase + bc;
+      let baseK = kBase + chunk * 16u;
+      var packed: u32 = 0u;
+      var scale: f32 = 0.0;
+      var zero: f32 = 0.0;
+      if (bn < N && baseK < K) {
+        let block = baseK / BLOCK_SIZE;
+        let offset = baseK % BLOCK_SIZE;
+        let byteIndex = (bn * KBLOCKS + block) * BLOB_SIZE + (offset >> 2u);
+        packed = b[byteIndex >> 2u];
+        scale = f32(scales[bn * KBLOCKS + block]);
+        zero = {{ zero_of("block") }};
+      }
+{% for vec in range(4) %}
+      let word{{ vec }} = vec4<f32>(
+{% for component in range(4) %}
+        (f32((packed >> {{ 2 * (vec * 4 + component) }}u) & 3u) - zero) * scale{% if not loop.last %},{% endif %}
+{% endfor %}
+      );
+      tileB[bc][chunk * 4u + {{ vec }}u] = word{{ vec }};
+{% endfor %}
+    }
+{% else %}
     // 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
       }
       tileB[bc][kv] = word;
     }
+{% endif %}
 {% else %}
     // Quant blocks that do not contain a whole BK tile use the element-wise loader.
     for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + WG_THREADS) {

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

@@ -42,18 +42,17 @@ fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
   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, kBlocks="KBLOCKS", blobSize="BLOB_SIZE") }}
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
-  let li = lid.y * 16u + lid.x;
   var acc00: f32 = 0.0;
   var acc01: f32 = 0.0;
@@ -65,7 +64,7 @@ fn main(
     // Cooperative load: 32x16 A tile + 16x32 B tile, 256 threads x 2 each. The B
     // tile is dequantized from the packed q blob during the load.
     for (var e: u32 = 0u; e < 2u; e = e + 1u) {
-      let idx = li + e * 256u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;

   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, kBlocks="KBLOCKS", blobSize="BLOB_SIZE") }}
+@compute @workgroup_size({{ tiledWorkgroupSide }}, {{ tiledWorkgroupSide }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>
 ) {
   let mBase = wg.y * BM;
   let nBase = wg.x * BN;
+  let li = lid.y * {{ tiledWorkgroupSide }}u + lid.x;
   var acc00: f32 = 0.0;
   var acc01: f32 = 0.0;
     // Cooperative load: 32x16 A tile + 16x32 B tile, 256 threads x 2 each. The B
     // tile is dequantized from the packed q blob during the load.
     for (var e: u32 = 0u; e < 2u; e = e + 1u) {
+      let idx = li + e * {{ tiledWorkgroupSide * tiledWorkgroupSide }}u;
       let ar = idx / BK;
       let ac = idx % BK;
       let am = mBase + ar;

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

@@ -13,7 +13,6 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 const M: u32 = {{ M }}u;
@@ -150,7 +149,7 @@ fn main(
     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" }});
@@ -169,6 +168,7 @@ fn main(
       matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
       matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
       matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
     }
     workgroupBarrier();
   }

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 const M: u32 = {{ M }}u;
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}{% set directInputs = false %}
       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" }});
       matC11 = subgroupMatrixMultiplyAccumulate(matA1, matB1, matC11);
       matC12 = subgroupMatrixMultiplyAccumulate(matA1, matB2, matC12);
       matC13 = subgroupMatrixMultiplyAccumulate(matA1, matB3, matC13);
     }
     workgroupBarrier();
   }

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

@@ -1,3 +1,8 @@
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
@@ -20,16 +25,12 @@ fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
   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) {

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const WG: u32 = {{ workgroupSize }}u;
   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>) {
+{{ flat_index_2d("WG", "index", "") }}
   let total = params.rows * params.N;
   if (index >= total) {

build/webgpu/metadata.json CHANGED Viewed

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

 {
   "name": "com.microsoft.MatMulNBits",
+  "id": "_com_microsoft_matmulnbits_webgpu_544d369",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "jsZTXUF4MjY52ewIINIKTdgXYZvvFfi+Ch6E3mCRIbo=",
+      "manifest.json": "C1+nZAAnjkvjvNZMOCbeF0zCAyPt3GIyqvbRUnz4KJo=",
+      "matmul-nbits-dp4a-quantize.wgsl.jinja": "7Ywrm3pRk4HfGeL5D0YcObG/kKLAgpr4dfu9vUJw8SM=",
+      "matmul-nbits-gemv-q4.wgsl.jinja": "0NkQ1G4eYRQGokjjMW7XKRvU76j8Dno9mrCBoXzeOc8=",
+      "matmul-nbits-q4-dp4a-prefill.wgsl.jinja": "w6qZuxXnJFteFjjDq+HonDn70RMHtPHreWxRtDiPIdU=",
       "matmul-nbits-q4-prefill-tile4x4.wgsl.jinja": "8DVy3szxcxwVItIlEXoioQi2YxDmmd5TCFW4BCYPiCQ=",
+      "matmul-nbits-q4-prefill-tiled-reg.wgsl.jinja": "Z22EQeNnROI5Vpy6mNgPK733TkxWXZB9rlp7qKV6fps=",
+      "matmul-nbits-q4-prefill-tiled.wgsl.jinja": "RzrALqTo5135awvgG96cJ029Hl8gsQtwujplaGYzcUs=",
+      "matmul-nbits-q4-sgmat.wgsl.jinja": "tUpNHvX4uEflOZ1HuFy4dezyL83kWO/5fYlxsHe5yaA=",
+      "matmul-nbits.wgsl.jinja": "i+euJpHwxE67Ejb8Pi9XzHDrK3FS/d2nT1odJJxR9rk=",
+      "reduce-axis0-splitk-combine.wgsl.jinja": "6S5tsaqhzGAZ66UOQ8u9LfKIlYeGnWu6B722kLl/auc=",
+      "test.json": "xR0ewiJSnSNGdmdIovbtC36Xsm1yv4aC8i3l8CaXdaY="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "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"],

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

@@ -2,31 +2,16 @@
 // 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;
@@ -48,7 +33,9 @@ fn is_nan_f32(value: f32) -> bool {
 @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" %}
@@ -74,13 +61,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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" %}
@@ -89,7 +69,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
     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];

 // 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 yv = "f16(" if outputF16 else "" %}
 {% set vy = ")" if outputF16 else "" %}
 {{ 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. */
 fn {{ name }}() -> {{ scalar }} {
   var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
   return bitcast<f32>(bits);
+}{% endmacro %}
 const WG: u32 = {{ workgroupSize }}u;
 const SPLIT: u32 = {{ split }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
+  // The start already folds gid.y in, so the stride must span every y row too;
+  // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own.
+  let stride = nwg.x * nwg.y * 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" %}
     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 op == "max" %}
     var total = reduction_identity();
 {% elif op == "min" %}
     var total = 1.0;
 {% else %}
     var total = 0.0;
 {% endif %}
     for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
       let p = partials[seg * params.cols + col];

build/webgpu/test.json CHANGED Viewed

@@ -4,7 +4,8 @@
     "q4_weight_cycle_b_t": [16, 50, 84, 118, 152, 186, 220, 254, 33, 67, 101, 135, 169, 203, 237, 15],
     "quant_scale_cycle_t": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07, 0.085],
     "mixed_weight_cycle_b_t": [17, 200, 91, 45, 233, 128, 7, 176, 250, 33, 142, 99, 210, 64, 188, 121],
-    "q8_zero_bias_gemv_m1_tail_n5_input_bT": [19, 56, 93, 130, 167, 204, 241, 22, 59, 96, 133, 170, 207, 244, 25, 62]
   },
   "cases": [
     {
@@ -42,7 +43,7 @@
     {
       "name": "q4_zero_bias_prefill_tile4x4_small_m8",
       "provenance": {
-        "notes": "Small-M (M=8) q4 prefill with bias and zero points. M<64 excludes prefill_tiled_zero_bias, while the row-guarded prefill_tile4x4_zero_bias route admits M>=2 when N is divisible by 4. This pins the tile4x4 lower-bound contract and its bias/zero-point arithmetic."
       },
       "inputs": {
         "aT": {
@@ -406,7 +407,7 @@
     {
       "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": {
@@ -440,9 +441,7 @@
     },
     {
       "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",
@@ -932,7 +931,7 @@
     {
       "name": "q4_no_zero_prefill_tile4x4_partial_row_tile_m6",
       "provenance": {
-        "notes": "M=6 is not a multiple of TILE_M=4, so the tile4x4 kernel's second row-tile (row_base=4) has valid rows 4,5 and guarded rows 6,7. Verifies the store_row partial-row-tile guard writes rows 4,5 correctly and does not corrupt/OOB rows 6,7. N=8 (%4==0), K=64, blockSize=32 routes to prefill_tile4x4_default_zero."
       },
       "inputs": {
         "aT": {
@@ -1033,7 +1032,7 @@
       "name": "q4_gemv_default_zero_m1_n13_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
-        "notes": "GEMV_N_COLS=8 with N=13: two workgroups, first fully live, second with a partially live first group and one live column in the second."
       },
       "inputs": {
         "aT": {
@@ -1059,7 +1058,7 @@
       "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": {
@@ -1184,7 +1183,7 @@
       "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": {
@@ -1219,7 +1218,7 @@
     {
       "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": {
@@ -1254,7 +1253,7 @@
     {
       "name": "q4_prefill_tiled_reg_tailk_m32_k33_n4096",
       "provenance": {
-        "notes": "Compact tail-block lock for the register-tiled prefill path used by the realistic K=2561 benchmark; the final 31 padded weights must not read past A."
       },
       "inputs": {
         "aT": {
@@ -1367,7 +1366,7 @@
     {
       "name": "q8_no_zero_prefill_odd_n_fallback",
       "provenance": {
-        "notes": "M>1 q8 prefill with N=17 and no zero_points/bias. Odd N excludes subgroup-matrix execution; the portable tile4x4 tail guards handle the final output column used by the odd-column benchmark guardrail."
       },
       "inputs": {
         "aT": {
@@ -2077,7 +2076,7 @@
     {
       "name": "q8_zero_only_naive_fallback_tailk_m3_n6",
       "provenance": {
-        "notes": "K=17 leaves a partial final block and M=3 bypasses GEMV. With explicit zero points but no bias, the aligned prefill paths are ineligible and the zero-only scalar fallback handles the tail block."
       },
       "inputs": {
         "aT": {
@@ -2107,7 +2106,7 @@
     {
       "name": "q8_bias_only_naive_fallback_tailk_m3_n6",
       "provenance": {
-        "notes": "K=17 leaves a partial final block and M=3 bypasses GEMV. With bias and the schema-default zero point, the aligned prefill paths are ineligible and the bias-only scalar fallback handles the tail block."
       },
       "inputs": {
         "aT": {
@@ -2278,6 +2277,177 @@
       },
       "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001 } },
       "attrs": { "K": 40, "N": 5, "bits": 4, "block_size": 16 }
     }
   ]
 }

     "q4_weight_cycle_b_t": [16, 50, 84, 118, 152, 186, 220, 254, 33, 67, 101, 135, 169, 203, 237, 15],
     "quant_scale_cycle_t": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07, 0.085],
     "mixed_weight_cycle_b_t": [17, 200, 91, 45, 233, 128, 7, 176, 250, 33, 142, 99, 210, 64, 188, 121],
+    "q8_zero_bias_gemv_m1_tail_n5_input_bT": [19, 56, 93, 130, 167, 204, 241, 22, 59, 96, 133, 170, 207, 244, 25, 62],
+    "ort_f16_large_k_accumulator_cancellation_m1_k8192_n8_input_scalesT": [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16, -16]
   },
   "cases": [
     {
     {
       "name": "q4_zero_bias_prefill_tile4x4_small_m8",
       "provenance": {
+        "notes": "M=8 rows of 4-bit quantized MatMulNBits with K=128, N=64, block_size=32, explicit zero points [6,7,8,9], and bias check correct dequantization and bias addition at a small row count."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_zero_bias_prefill_reg_vec4_splitk_m128_k1024_n1024",
       "provenance": {
+        "notes": "Zero points affect the complete reduction, and bias is added exactly once to the final output."
       },
       "inputs": {
         "aT": {
     },
     {
       "name": "q8_zero_only_prefill_reg_vec4_splitk_m128_k1024_n1024",
+      "provenance": { "notes": "Zero points affect the complete reduction without an output bias." },
       "inputs": {
         "aT": {
           "dtype": "float32",
     {
       "name": "q4_no_zero_prefill_tile4x4_partial_row_tile_m6",
       "provenance": {
+        "notes": "M=6 rows of 4-bit MatMulNBits (K=64, N=8, block_size=32, default zero point, no bias) check a row count that is not a multiple of four."
       },
       "inputs": {
         "aT": {
       "name": "q4_gemv_default_zero_m1_n13_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
+        "notes": "N=13 with M=1, K=32, block_size=32, 4-bit weights and the default zero point checks a column count that is not a multiple of eight."
       },
       "inputs": {
         "aT": {
       "name": "q4_gemv_default_zero_m1_tail_n7_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
+        "notes": "N=7 with M=1, K=32, block_size=32, 4-bit weights and the default zero point checks that all seven output columns are dequantized and written correctly."
       },
       "inputs": {
         "aT": {
       "name": "q8_zero_bias_gemv_m1_tail_n5_ncols8",
       "tunables": { "GEMV_N_COLS": 8 },
       "provenance": {
+        "notes": "N=5 with M=1, K=16, block_size=16, 8-bit weights, explicit per-block zero points and a bias checks that all five output columns are dequantized and written correctly."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_zero_bias_naive_fallback_tailK_m3_n6",
       "provenance": {
+        "notes": "K=17 with block size 16 leaves a partial quantization block. M=3 and N=6 check q8 unpacking, per-block zero points, bias, and the K tail."
       },
       "inputs": {
         "aT": {
     {
       "name": "q4_prefill_tiled_reg_tailk_m32_k33_n4096",
       "provenance": {
+        "notes": "A 33-element reduction leaves a partial final quantization block; padded weights must not contribute to the result or read past the activation input."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_no_zero_prefill_odd_n_fallback",
       "provenance": {
+        "notes": "An eight-row q8 projection with N=17 and no zero points or bias checks the final odd output column."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_zero_only_naive_fallback_tailk_m3_n6",
       "provenance": {
+        "notes": "K=17 with block_size=16 leaves one element in a partial final quantization block (M=3, N=6, 8-bit, explicit per-block zero points, no bias); checks that the partial block is weighted correctly without bias."
       },
       "inputs": {
         "aT": {
     {
       "name": "q8_bias_only_naive_fallback_tailk_m3_n6",
       "provenance": {
+        "notes": "K=17 with block_size=16 leaves one element in a partial final quantization block (M=3, N=6, 8-bit, default zero point, explicit bias); checks that the partial block and bias combine correctly."
       },
       "inputs": {
         "aT": {
       },
       "outputs": { "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001 } },
       "attrs": { "K": 40, "N": 5, "bits": 4, "block_size": 16 }
+    },
+    {
+      "name": "ort_f16_large_k_accumulator_cancellation_m1_k8192_n8",
+      "provenance": {
+        "notes": "Transcribed from ORT MatMulNBits.Float16_LargeK_AccumulatorOverflow (M=1 arm). A = 8 everywhere; the dequantized weight is +112 over the first half of K and -112 over the second, so the exact result is 0 while the running partial sum crosses the f16 ceiling (65504) at 114688 in any kernel that walks K in one accumulator. ORT puts the sign on B's codes (0xFF then 0x11); this package puts it on the block scale instead (+16 then -16) with every code 15, which yields a bit-identical dequantized weight matrix and lets B be one constant instead of a 4096-long cycle."
+      },
+      "attrs": { "K": 8192, "N": 8, "bits": 4, "block_size": 32 },
+      "inputs": {
+        "aT": { "dtype": "float16", "shape": [1, 8192], "data": { "kind": "constant", "value": 8.0 } },
+        "bT": { "dtype": "uint8", "shape": [8, 256, 16], "data": { "kind": "constant", "value": 255 } },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [8, 256],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/ort_f16_large_k_accumulator_cancellation_m1_k8192_n8_input_scalesT" }
+          }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [1, 8], "tolerance": 0.05 } }
+    },
+    {
+      "name": "ort_f16_large_k_accumulator_cancellation_m8_k8192_n8",
+      "provenance": {
+        "notes": "Transcribed from ORT MatMulNBits.Float16_LargeK_AccumulatorOverflow (M=8 arm). The same cancellation construction as the M=1 case checks float32 accumulation for eight output rows."
+      },
+      "attrs": { "K": 8192, "N": 8, "bits": 4, "block_size": 32 },
+      "inputs": {
+        "aT": { "dtype": "float16", "shape": [8, 8192], "data": { "kind": "constant", "value": 8.0 } },
+        "bT": { "dtype": "uint8", "shape": [8, 256, 16], "data": { "kind": "constant", "value": 255 } },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [8, 256],
+          "data": {
+            "kind": "cycle",
+            "values": { "$ref": "#/fixtureArrays/ort_f16_large_k_accumulator_cancellation_m1_k8192_n8_input_scalesT" }
+          }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [8, 8], "tolerance": 0.05 } }
+    },
+    {
+      "name": "f16_full_k_accumulator_cancellation_prefill_m64_k4096_n64",
+      "provenance": {
+        "notes": "Adapted from ORT MatMulNBits.Float16_LargeK_AccumulatorOverflow. With K=4,096, partial sums rise above the float16 finite ceiling (peak 1,835,008) before cancelling to an exact zero output."
+      },
+      "attrs": { "K": 4096, "N": 64, "bits": 4, "block_size": 32 },
+      "inputs": {
+        "aT": { "dtype": "float16", "shape": [64, 4096], "data": { "kind": "constant", "value": 8.0 } },
+        "bT": { "dtype": "uint8", "shape": [64, 128, 16], "data": { "kind": "constant", "value": 255 } },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [64, 128],
+          "data": {
+            "kind": "cycle",
+            "values": [16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0, -16.0]
+          }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [64, 64], "tolerance": 0.05 } }
+    },
+    {
+      "name": "ort_f16_decode_bias_m1_k1024_n128_b32",
+      "provenance": {
+        "notes": "Transcribed from ORT MatMulNBits.Float16_AccumulatorPrecisionOption_AllPaths, case {M=1, N=128, K=1024, block 32, accuracy_level 0, bias}. The provider option is not applicable here (accumulators are always f32); the shape is kept because it is the generic decode dispatch with a bias and there was no other float16 M=1 case. Weight/scale data reuse this file's existing fixture arrays."
+      },
+      "attrs": { "K": 1024, "N": 128, "bits": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [1, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.25 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [128, 32, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [128, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "biasT": {
+          "dtype": "float16",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.125 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [1, 128], "tolerance": 0.01, "relTolerance": 0.005 } }
+    },
+    {
+      "name": "ort_f16_prefill_bias_m8_k1024_n128_b32",
+      "provenance": {
+        "notes": "Transcribed from ORT MatMulNBits.Float16_AccumulatorPrecisionOption_AllPaths, case {M=8, N=128, K=1024, block 32, accuracy_level 0, bias} (ORT's wide-tile arm)."
+      },
+      "attrs": { "K": 1024, "N": 128, "bits": 4, "block_size": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [8, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.25 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [128, 32, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [128, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "biasT": {
+          "dtype": "float16",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.125 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [8, 128], "tolerance": 0.01, "relTolerance": 0.005 } }
+    },
+    {
+      "name": "q2_block128_f16_zero1_prefill_m128_k1024_n1024",
+      "provenance": {
+        "notes": "M=128, K=1024, N=1024 with 2-bit weights in 128-value quantization blocks, a uniform explicit zero point of 1, and float16 operands check dequantization at a full-size projection shape."
+      },
+      "attrs": { "K": 1024, "N": 1024, "bits": 2, "block_size": 128 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [128, 1024],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.25 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024, 8, 32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/mixed_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [1024, 8],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        },
+        "zeroPointsT": { "dtype": "float16", "shape": [1024, 8], "data": { "kind": "constant", "value": 1.0 } }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [128, 1024], "tolerance": 0.01, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "ort_f16_accuracy_level4_m8_k4096_n128_generic_route",
+      "provenance": {
+        "notes": "Transcribed from ORT MatMulNBits.Float16_AccumulatorPrecisionOption_AllPaths: M=8, N=128, K=4,096, block size 32, accuracy_level 4. The float16 inputs check float32 accumulation independently of provider-specific execution choices."
+      },
+      "attrs": { "K": 4096, "N": 128, "bits": 4, "block_size": 32, "accuracy_level": 4 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [8, 4096],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23, "scale": 0.25 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [128, 128, 16],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/q4_weight_cycle_b_t" } }
+        },
+        "scalesT": {
+          "dtype": "float16",
+          "shape": [128, 128],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/quant_scale_cycle_t" } }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [8, 128], "tolerance": 0.02, "relTolerance": 0.01 } }
     }
   ]
 }