sync 2e7068faf55e

Browse files

Files changed (10) hide show

README.md +91 -0
build/webgpu/bench.json +122 -0
build/webgpu/cast-scalar-x4.wgsl.jinja +27 -0
build/webgpu/manifest.json +390 -0
build/webgpu/matmul-bnb4-gemv.wgsl.jinja +88 -0
build/webgpu/matmul-bnb4-sgmat.wgsl.jinja +128 -0
build/webgpu/matmul-bnb4-tiled.wgsl.jinja +168 -0
build/webgpu/matmul-bnb4.wgsl.jinja +65 -0
build/webgpu/metadata.json +22 -0
build/webgpu/test.json +434 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.MatMulBnb4
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Computes `A @ dequant(B)^T` where `B` uses bitsandbytes 4-bit quantization: `quant_type = 0` selects FP4 and `quant_type = 1` selects NF4. Supports rank-2 float16/float32 `A`, `transB = 1`, and `training_mode = 0`; rank-1 and rank-3-or-higher `A`, bfloat16, `transB = 0`, and training are not implemented. `B` is the flattened `[N, K]` weight, two codes per byte with the even flat index in the high nibble. Each code indexes a fixed 16-entry codebook, and the value is `codebook[code] * absmax[flat_index / block_size]`.
+See the [ONNX Runtime `MatMulBnb4` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MatMulBnb4) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `A` | `aT` | `T1` | `2` | — | Float input matrix of shape `(M, K)`, not quantized. | required |
+| `B` | `bT` | `T2` | `1` | — | The `[N, K]` weight, flattened and quantized to 4 bits, stored as `(N * K + 1) / 2` bytes; the ONNX type is uint8 (this WebGPU implementation reads one widened u32 per stored byte). | required |
+| `absmax` | `absmaxT` | `T1` | `1` | — | Per-block absolute-maximum dequantization scales of shape `((N * K + block_size - 1) / block_size)`, same dtype as A. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `yT` | `T1` | `2` | `[A[0], N]` | Result of `A` multiplied by the dequantized, transposed weight matrix, with shape `(M, N)` and the same dtype as `A`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `training_mode` | `0` | Whether training outputs are requested. This inference-only implementation supports the standard default value 0. |
+| `transB` | `1` | Whether the quantized weight is stored transposed. This implementation supports the standard default value 1. |
+| `K` | — | Input feature count (the shared dimension). |
+| `N` | — | Output feature count. |
+| `block_size` | — | Number of weights sharing one absmax scale; a power of two, at least 16. |
+| `quant_type` | — | Codebook selector: 0 = FP4, 1 = NF4. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T1` | `float32`, `float16` |
+| `T2` | `uint8` |
+## Device requirements
+Some implementation variants require `subgroup-matrix` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`cast-scalar-x4.wgsl.jinja`](build/webgpu/cast-scalar-x4.wgsl.jinja)
+- [`matmul-bnb4-gemv.wgsl.jinja`](build/webgpu/matmul-bnb4-gemv.wgsl.jinja)
+- [`matmul-bnb4-sgmat.wgsl.jinja`](build/webgpu/matmul-bnb4-sgmat.wgsl.jinja)
+- [`matmul-bnb4-tiled.wgsl.jinja`](build/webgpu/matmul-bnb4-tiled.wgsl.jinja)
+- [`matmul-bnb4.wgsl.jinja`](build/webgpu/matmul-bnb4.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.MatMulBnb4", { version: 1 });
+const { yT } = await kernel({
+  aT: { data: aTData, shape: [2, 24] },
+  bT: { data: bTData, shape: [36] },
+  absmaxT: { data: absmaxTData, shape: [5] },
+}, {
+  attrs: {
+    K: 24,
+    N: 3,
+    block_size: 16,
+    quant_type: 1,
+  },
+});
+```

build/webgpu/bench.json ADDED Viewed

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

+{
+  "op": "com.microsoft.MatMulBnb4",
+  "cases": [
+    {
+      "name": "bnb4-nf4-decode-m1-k2048-n2048-b64",
+      "preset": "smoke",
+      "vars": { "M": 1, "K": 2048, "N": 2048, "blockSize": 64 },
+      "attrs": { "K": 2048, "N": 2048, "block_size": 64, "quant_type": 1 },
+      "inputs": {
+        "aT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 410, "scale": 0.2 },
+        "bT": { "shape": [2097152], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "absmaxT": {
+          "shape": [65536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 411,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [1, 2048], "dtype": "float32" } },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "numel(shapes.bT) * 4 + (numel(shapes.aT) + numel(shapes.absmaxT) + numel(shapes.yT)) * 4"
+          }
+        ]
+      }
+    },
+    {
+      "name": "bnb4-nf4-prefill-m128-k2048-n2048-b64",
+      "preset": "smoke",
+      "vars": { "M": 128, "K": 2048, "N": 2048, "blockSize": 64 },
+      "attrs": { "K": 2048, "N": 2048, "block_size": 64, "quant_type": 1 },
+      "tunableSpace": { "PORTABLE_TILE_K": [16, 32, 64] },
+      "inputs": {
+        "aT": { "shape": [128, 2048], "dtype": "float32", "dist": "normal", "seed": 412, "scale": 0.2 },
+        "bT": { "shape": [2097152], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "absmaxT": {
+          "shape": [65536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 413,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [128, 2048], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "gflops", "value": "2 * dim(shapes.aT, 0) * args.K * args.N" }]
+      }
+    },
+    {
+      "name": "bnb4-fp4-prefill-m128-k2048-n2048-b64",
+      "preset": "smoke",
+      "vars": { "M": 128, "K": 2048, "N": 2048, "blockSize": 64 },
+      "attrs": { "K": 2048, "N": 2048, "block_size": 64, "quant_type": 0 },
+      "tunableSpace": { "PORTABLE_TILE_K": [16, 32, 64] },
+      "inputs": {
+        "aT": { "shape": [128, 2048], "dtype": "float32", "dist": "normal", "seed": 414, "scale": 0.2 },
+        "bT": { "shape": [2097152], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "absmaxT": {
+          "shape": [65536],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 415,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [128, 2048], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * dim(shapes.aT, 0) * args.K * args.N" }] }
+    },
+    {
+      "name": "bnb4-nf4-prefill-m512-k4096-n4096-b64",
+      "preset": "model",
+      "vars": { "M": 512, "K": 4096, "N": 4096, "blockSize": 64 },
+      "attrs": { "K": 4096, "N": 4096, "block_size": 64, "quant_type": 1 },
+      "inputs": {
+        "aT": { "shape": [512, 4096], "dtype": "float32", "dist": "normal", "seed": 416, "scale": 0.2 },
+        "bT": { "shape": [8388608], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "absmaxT": {
+          "shape": [262144],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 417,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [512, 4096], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * dim(shapes.aT, 0) * args.K * args.N" }] }
+    },
+    {
+      "name": "bnb4-nf4-prefill-m512-k4096-n4096-b64-f16",
+      "preset": "model",
+      "vars": { "M": 512, "K": 4096, "N": 4096, "blockSize": 64 },
+      "attrs": { "K": 4096, "N": 4096, "block_size": 64, "quant_type": 1 },
+      "inputs": {
+        "aT": { "shape": [512, 4096], "dtype": "float16", "dist": "normal", "seed": 416, "scale": 0.2 },
+        "bT": { "shape": [8388608], "dtype": "uint8", "dist": "linearMod", "mod": 256, "step": 7 },
+        "absmaxT": {
+          "shape": [262144],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 417,
+          "offset": 0.04,
+          "scale": 0.01,
+          "signed": false
+        }
+      },
+      "outputs": { "yT": { "shape": [512, 4096], "dtype": "float16" } },
+      "bench": { "metrics": [{ "type": "gflops", "value": "2 * dim(shapes.aT, 0) * args.K * args.N" }] }
+    }
+  ]
+}

build/webgpu/cast-scalar-x4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,27 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% if wrapNarrowInt %}
+// ONNX float-to-int8/uint8 narrowing wraps modulo 256 rather than saturating.
+// Preserve the low byte after truncation; the signed path then sign-extends it.
+fn cast_wrap(v: f32) -> {{ outScalar }} {
+  let low = i32(v) & 0xFF;
+  return {{ outScalar }}({% if wrapSigned %}select(low, low - 256, low > 127){% else %}low{% endif %});
+}
+{% endif %}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // Four scalar conversions per invocation retain vector-like dispatch density even when the
+  // logical element count is not vec4 aligned. For a vec4 bulk + scalar tail plan, only lane zero
+  // of the tail dispatch runs and starts at the first element not covered by the packed pass.
+  let invocation = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  let base = invocation * 4u;
+{% for lane in range(4) %}
+  if (base + {{ lane }}u < params.count) {
+    y[base + {{ lane }}u] = {% if wrapNarrowInt %}cast_wrap(f32(x[base + {{ lane }}u])){% else %}{{ outScalar }}(x[base + {{ lane }}u]){% endif %};
+  }
+{% endfor %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,390 @@

+{
+  "domain": "com.microsoft",
+  "name": "MatMulBnb4",
+  "sinceVersion": 1,
+  "description": "Computes `A @ dequant(B)^T` where `B` uses bitsandbytes 4-bit quantization: `quant_type = 0` selects FP4 and `quant_type = 1` selects NF4. Supports rank-2 float16/float32 `A`, `transB = 1`, and `training_mode = 0`; rank-1 and rank-3-or-higher `A`, bfloat16, `transB = 0`, and training are not implemented. `B` is the flattened `[N, K]` weight, two codes per byte with the even flat index in the high nibble. Each code indexes a fixed 16-entry codebook, and the value is `codebook[code] * absmax[flat_index / block_size]`.",
+  "inputs": [
+    { "role": "A", "dtype": "T1", "rank": 2, "description": "Float input matrix of shape `(M, K)`, not quantized." },
+    {
+      "role": "B",
+      "dtype": "T2",
+      "rank": 1,
+      "description": "The `[N, K]` weight, flattened and quantized to 4 bits, stored as `(N * K + 1) / 2` bytes; the ONNX type is uint8 (this WebGPU implementation reads one widened u32 per stored byte)."
+    },
+    {
+      "role": "absmax",
+      "dtype": "T1",
+      "rank": 1,
+      "description": "Per-block absolute-maximum dequantization scales of shape `((N * K + block_size - 1) / block_size)`, same dtype as A."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "Y",
+      "dtype": "T1",
+      "rank": 2,
+      "shape": "[dim(shapes.A, 0), attrs.N]",
+      "description": "Result of `A` multiplied by the dequantized, transposed weight matrix, with shape `(M, N)` and the same dtype as `A`."
+    }
+  ],
+  "attributes": { "training_mode": 0, "transB": 1 },
+  "attributeConstraints": {
+    "K": { "required": true },
+    "N": { "required": true },
+    "block_size": { "required": true },
+    "quant_type": { "required": true, "values": [0, 1] },
+    "training_mode": { "values": [0] },
+    "transB": { "values": [1] }
+  },
+  "attributeDescriptions": {
+    "K": "Input feature count (the shared dimension).",
+    "N": "Output feature count.",
+    "block_size": "Number of weights sharing one absmax scale; a power of two, at least 16.",
+    "quant_type": "Codebook selector: 0 = FP4, 1 = NF4.",
+    "training_mode": "Whether training outputs are requested. This inference-only implementation supports the standard default value 0.",
+    "transB": "Whether the quantized weight is stored transposed. This implementation supports the standard default value 1."
+  },
+  "typeConstraints": { "T1": ["float32", "float16"], "T2": ["uint8"] },
+  "args": {
+    "aT": { "kind": "tensor", "semantic": "A", "role": "input" },
+    "bT": { "kind": "tensor", "semantic": "B", "role": "input" },
+    "absmaxT": { "kind": "tensor", "semantic": "absmax", "role": "input" },
+    "yT": { "kind": "tensor", "semantic": "Y", "role": "output" }
+  },
+  "tunables": {
+    "WORKGROUP_SIZE": 64,
+    "TILE_MIN_M": 16,
+    "PORTABLE_TILE_K": 16,
+    "SGMAT_TILE_ROWS": 64,
+    "SGMAT_TALL_TILE_ROWS": 128,
+    "SGMAT_TALL_MIN_M": 128,
+    "SGMAT_TILE_COLS": 64,
+    "SGMAT_TILE_K": 32
+  },
+  "bindingSets": {
+    "main": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "absmax",
+        "arg": "absmaxT",
+        "semantic": "absmax",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$absmaxScalar",
+        "length": "$ABSMAX_LEN"
+      },
+      { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "dim(shapes.A, 0)" },
+            { "name": "K", "type": "u32", "value": "attrs.K" },
+            { "name": "N", "type": "u32", "value": "attrs.N" },
+            { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
+          ]
+        }
+      }
+    ],
+    "gemv": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "absmax",
+        "arg": "absmaxT",
+        "semantic": "absmax",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$absmaxScalar",
+        "length": "$ABSMAX_LEN"
+      },
+      { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "K", "type": "u32", "value": "attrs.K" },
+            { "name": "N", "type": "u32", "value": "attrs.N" },
+            { "name": "blockSize", "type": "u32", "value": "attrs.block_size" }
+          ]
+        }
+      }
+    ],
+    "sgmat": [
+      {
+        "name": "a",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$aScalar"
+      },
+      { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "absmax",
+        "arg": "absmaxT",
+        "semantic": "absmax",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$absmaxScalar",
+        "length": "$ABSMAX_LEN"
+      },
+      { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$aScalar" }
+    ],
+    "castA": [
+      {
+        "name": "x",
+        "arg": "aT",
+        "semantic": "A",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$srcScalar"
+      },
+      { "name": "y", "semantic": "aF32", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.A)" }] }
+      }
+    ],
+    "sgmatWiden": [
+      { "name": "a", "semantic": "aF32", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "b", "arg": "bT", "semantic": "B", "buffer": { "type": "read-only-storage" }, "elementType": "u32" },
+      {
+        "name": "absmax",
+        "arg": "absmaxT",
+        "semantic": "absmax",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$absmaxScalar",
+        "length": "$ABSMAX_LEN"
+      },
+      { "name": "y", "semantic": "yF32", "buffer": { "type": "storage" }, "elementType": "f32" }
+    ],
+    "castY": [
+      { "name": "x", "semantic": "yF32", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "y", "arg": "yT", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$srcScalar" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [{ "name": "count", "type": "u32", "value": "dim(shapes.A, 0) * attrs.N" }]
+        }
+      }
+    ]
+  },
+  "derive": {
+    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "packedBytesExpected": "ceilDiv(attrs.N * attrs.K, 2)",
+    "absmaxCountExpected": "ceilDiv(attrs.N * attrs.K, attrs.block_size)",
+    "aFloatOk": "(tensorDtypes.A == \"float32\" or tensorDtypes.A == \"float16\") and f16Ok(tensorDtypes.A)",
+    "portableWorkgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "commonShapeValid": "ranks.A == 2 and ranks.B == 1 and ranks.absmax == 1 and ranks.Y == 2 and aFloatOk and tensorDtypes.B == \"uint8\" and tensorDtypes.absmax == tensorDtypes.A and tensorDtypes.Y == tensorDtypes.A and attrs.K > 0 and attrs.N > 0 and attrs.block_size >= 16 and pow2ceil(attrs.block_size) == attrs.block_size and dim(shapes.A, 1) == attrs.K and dim(shapes.B, 0) == packedBytesExpected and dim(shapes.absmax, 0) == absmaxCountExpected and dim(shapes.Y, 0) == dim(shapes.A, 0) and dim(shapes.Y, 1) == attrs.N",
+    "gemvShapeValid": "commonShapeValid and dim(shapes.A, 0) == 1",
+    "portableWorkgroupFits": "portableWorkgroupSize > 0 and portableWorkgroupSize * 16 <= device.limits.maxComputeWorkgroupStorageSize",
+    "portableTileKValid": "tunables.PORTABLE_TILE_K >= 8 and tunables.PORTABLE_TILE_K % 8 == 0",
+    "tileEligible": "commonShapeValid and dim(shapes.A, 0) >= tunables.TILE_MIN_M and portableTileKValid",
+    "tileWorkgroupStorageBytes": "8 * 64 * tunables.PORTABLE_TILE_K",
+    "tileWorkgroupFits": "16 <= device.limits.maxComputeWorkgroupSizeX and 8 <= device.limits.maxComputeWorkgroupSizeY and 128 <= device.limits.maxComputeInvocationsPerWorkgroup and tileWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(attrs.N, 64) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.A, 0), 64) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "sgmatMatrixSize": "8",
+    "sgmatTileRows": "tunables.SGMAT_TALL_TILE_ROWS if dim(shapes.A, 0) >= tunables.SGMAT_TALL_MIN_M else tunables.SGMAT_TILE_ROWS",
+    "sgmatRowSubtiles": "4",
+    "sgmatSubRows": "sgmatTileRows / sgmatRowSubtiles",
+    "sgmatSubCols": "4 * sgmatMatrixSize",
+    "sgmatLoadWidth": "sgmatMatrixSize",
+    "sgmatColSubtiles": "tunables.SGMAT_TILE_COLS / sgmatSubCols",
+    "sgmatNumSubgroups": "sgmatRowSubtiles * sgmatColSubtiles",
+    "sgmatSubgroupSize": "device.adapterInfo.subgroupMinSize if has(device.adapterInfo, \"subgroupMinSize\") else 1",
+    "sgmatWorkgroupSize": "sgmatNumSubgroups * sgmatSubgroupSize",
+    "sgmatBLoadsPerRow": "tunables.SGMAT_TILE_K / sgmatLoadWidth",
+    "sgmatWorkgroupStorageBytes": "4 * tunables.SGMAT_TILE_COLS * tunables.SGMAT_TILE_K",
+    "sgmatDispatchN": "ceilDiv(attrs.N, tunables.SGMAT_TILE_COLS)",
+    "sgmatDispatchM": "ceilDiv(dim(shapes.A, 0), sgmatTileRows)",
+    "sgmatWorkgroupFits": "sgmatWorkgroupSize <= deviceWorkgroupCap and sgmatWorkgroupStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and sgmatDispatchN <= device.limits.maxComputeWorkgroupsPerDimension and sgmatDispatchM <= device.limits.maxComputeWorkgroupsPerDimension",
+    "sgmatWidenBytes": "numel(shapes.A) * 4",
+    "sgmatWidenOutBytes": "dim(shapes.A, 0) * attrs.N * 4",
+    "sgmatWidenFits": "sgmatWidenBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenBytes <= device.limits.maxBufferSize and sgmatWidenOutBytes <= device.limits.maxStorageBufferBindingSize and sgmatWidenOutBytes <= device.limits.maxBufferSize"
+  },
+  "constants": { "quantType": "attrs.quant_type", "ABSMAX_LEN": "ceilDiv(attrs.N * attrs.K, attrs.block_size)" },
+  "variants": [
+    {
+      "id": "gemv",
+      "priority": 20,
+      "when": ["gemvShapeValid", "portableWorkgroupFits"],
+      "constants": {
+        "workgroupSize": "portableWorkgroupSize",
+        "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"",
+        "usesF16": "tensorDtypes.A == \"float16\"",
+        "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "matmul-bnb4-gemv.wgsl.jinja",
+          "bindings": "gemv",
+          "dispatch": { "workgroups": "attrs.N" }
+        }
+      ]
+    },
+    {
+      "id": "sgmat",
+      "priority": 15,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "limits": { "maxComputeWorkgroupStorageSize": 8192 },
+        "subgroupMinSize": 32,
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "when": ["commonShapeValid", "tensorDtypes.A == \"float32\"", "dim(shapes.A, 0) >= sgmatTileRows", "dim(shapes.A, 0) % sgmatTileRows == 0", "attrs.N % tunables.SGMAT_TILE_COLS == 0", "attrs.K % tunables.SGMAT_TILE_K == 0", "attrs.K % attrs.block_size == 0", "attrs.block_size % sgmatLoadWidth == 0", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "sgmatWorkgroupFits"],
+      "constants": {
+        "K": "attrs.K",
+        "N": "attrs.N",
+        "blockSize": "attrs.block_size",
+        "tileRows": "sgmatTileRows",
+        "tileCols": "tunables.SGMAT_TILE_COLS",
+        "tileK": "tunables.SGMAT_TILE_K",
+        "subRows": "sgmatSubRows",
+        "subCols": "sgmatSubCols",
+        "matrixSize": "sgmatMatrixSize",
+        "rowMatrices": "sgmatSubRows / sgmatMatrixSize",
+        "colMatrices": "sgmatSubCols / sgmatMatrixSize",
+        "loadWidth": "sgmatLoadWidth",
+        "rowSubtiles": "sgmatRowSubtiles",
+        "workgroupSize": "sgmatWorkgroupSize",
+        "bLoadsPerRow": "sgmatBLoadsPerRow",
+        "aScalar": "\"f32\"",
+        "usesF16": false,
+        "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "matmul-bnb4-sgmat.wgsl.jinja",
+          "bindings": "sgmat",
+          "dispatch": { "x": "sgmatDispatchN", "y": "sgmatDispatchM" }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_widened",
+      "priority": 16,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "limits": { "maxComputeWorkgroupStorageSize": 8192 },
+        "subgroupMinSize": 32,
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "when": ["commonShapeValid", "tensorDtypes.A == \"float16\"", "dim(shapes.A, 0) >= sgmatTileRows", "dim(shapes.A, 0) % sgmatTileRows == 0", "attrs.N % tunables.SGMAT_TILE_COLS == 0", "attrs.K % tunables.SGMAT_TILE_K == 0", "attrs.K % attrs.block_size == 0", "attrs.block_size % sgmatLoadWidth == 0", "has(device.adapterInfo, \"subgroupMinSize\")", "has(device.adapterInfo, \"subgroupMaxSize\")", "device.adapterInfo.subgroupMinSize == 32", "device.adapterInfo.subgroupMaxSize == 32", "sgmatWorkgroupFits", "sgmatWidenFits"],
+      "constants": {
+        "K": "attrs.K",
+        "N": "attrs.N",
+        "blockSize": "attrs.block_size",
+        "tileRows": "sgmatTileRows",
+        "tileCols": "tunables.SGMAT_TILE_COLS",
+        "tileK": "tunables.SGMAT_TILE_K",
+        "subRows": "sgmatSubRows",
+        "subCols": "sgmatSubCols",
+        "matrixSize": "sgmatMatrixSize",
+        "rowMatrices": "sgmatSubRows / sgmatMatrixSize",
+        "colMatrices": "sgmatSubCols / sgmatMatrixSize",
+        "loadWidth": "sgmatLoadWidth",
+        "rowSubtiles": "sgmatRowSubtiles",
+        "workgroupSize": "sgmatWorkgroupSize",
+        "bLoadsPerRow": "sgmatBLoadsPerRow",
+        "aScalar": "\"f32\"",
+        "usesF16": true,
+        "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\"",
+        "srcScalar": "\"f16\"",
+        "outScalar": "\"f32\"",
+        "wrapNarrowInt": false,
+        "wrapSigned": false
+      },
+      "passes": [
+        {
+          "id": "widen_a",
+          "name": "MatMulBnb4.WidenActivations",
+          "shader": "cast-scalar-x4.wgsl.jinja",
+          "bindings": "castA",
+          "dispatch": { "threads": "ceilDiv(numel(shapes.A), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        },
+        {
+          "id": "main",
+          "name": "MatMulBnb4.SubgroupMatrixWidened",
+          "shader": "matmul-bnb4-sgmat.wgsl.jinja",
+          "bindings": "sgmatWiden",
+          "dispatch": { "x": "sgmatDispatchN", "y": "sgmatDispatchM" }
+        },
+        {
+          "id": "narrow_y",
+          "name": "MatMulBnb4.NarrowOutput",
+          "shader": "cast-scalar-x4.wgsl.jinja",
+          "bindings": "castY",
+          "constants": { "outScalar": "\"f16\"" },
+          "dispatch": {
+            "threads": "ceilDiv(dim(shapes.A, 0) * attrs.N, 4)",
+            "workgroupSize": "tunables.WORKGROUP_SIZE"
+          }
+        }
+      ],
+      "description": "Runs the subgroup-matrix tier for a float16 request by widening the activations to float32 either side of the multiply. The matrix units this operator uses accumulate in float32 from float32 operands; the device's float16 configuration returns a float16 result, so feeding them float16 directly would drop the accumulator's precision.",
+      "intermediates": [
+        { "id": "aF32", "dtype": "float32", "shape": "[numel(shapes.A)]" },
+        { "id": "yF32", "dtype": "float32", "shape": "[dim(shapes.A, 0) * attrs.N]" }
+      ]
+    },
+    {
+      "id": "tiled",
+      "priority": 10,
+      "when": ["tileEligible", "tileWorkgroupFits"],
+      "constants": {
+        "tileK": "tunables.PORTABLE_TILE_K",
+        "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"",
+        "usesF16": "tensorDtypes.A == \"float16\"",
+        "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "matmul-bnb4-tiled.wgsl.jinja",
+          "bindings": "main",
+          "dispatch": { "x": "ceilDiv(attrs.N, 64)", "y": "ceilDiv(dim(shapes.A, 0), 64)" }
+        }
+      ]
+    },
+    {
+      "id": "scalar",
+      "priority": 0,
+      "when": ["commonShapeValid", "portableWorkgroupSize > 0"],
+      "constants": {
+        "workgroupSize": "portableWorkgroupSize",
+        "aScalar": "\"f16\" if tensorDtypes.A == \"float16\" else \"f32\"",
+        "usesF16": "tensorDtypes.A == \"float16\"",
+        "absmaxScalar": "\"f16\" if tensorDtypes.absmax == \"float16\" else \"f32\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "shader": "matmul-bnb4.wgsl.jinja",
+          "bindings": "main",
+          "dispatch": { "threads": "numel(shapes.Y)", "workgroupSize": "constants.workgroupSize" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/matmul-bnb4-gemv.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,88 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+var<workgroup> partials: array<f32, {{ workgroupSize }}>;
+/* bitsandbytes 4-bit weight decoding.
+   Only the selected FP4 or NF4 codebook is emitted. The codebook uses
+   var<private> rather than const because its runtime index requires a memory
+   location in WGSL.
+   FP4 is sign-magnitude (bit 3 is sign), so its table is non-monotonic. NF4 is
+   the ascending normal-float quantile codebook. Each stored byte is carried in
+   one widened u32. A flat weight index maps to byte flat >> 1, and an even
+   index reads the high nibble, matching (vi0 << 4) | vi1 packing. */
+{% if quantType == 0 %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  0.0, 5.208333333e-03, 0.66666667, 1.0,
+  0.33333333, 0.5, 0.16666667, 0.25,
+  -0.0, -5.208333333e-03, -0.66666667, -1.0,
+  -0.33333333, -0.5, -0.16666667, -0.25
+);
+{%- else %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
+  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
+  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
+  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
+);
+{%- endif %}
+// One dequantized weight at flat index `flat` into the [N, K] weight.
+fn dequant_weight(flat: u32) -> f32 {
+  let byte = b[flat >> 1u];
+  let code = select(byte & 0xFu, (byte >> 4u) & 0xFu, (flat & 1u) == 0u);
+  return quant_map[code] * f32(absmax[flat / params.blockSize]);
+}
+// Decode (M == 1): one workgroup owns each output column, with lanes striding K.
+// Each lane keeps one accumulator and the workgroup closes with a tree reduction.
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(workgroup_id) wid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let tid = lid.x;
+  // 2D-folded dispatch: wid.y carries the high bits past the per-dimension
+  // limit. The fold over-dispatches, so tail workgroups must leave before the
+  // first barrier — hence the early return here rather than a guarded write.
+  let col = wid.x + wid.y * nwg.x;
+  if (col >= params.N) {
+    return;
+  }
+  var acc = 0.0;
+  let b_row_base = col * params.K;
+  for (var k = tid; k < params.K; k = k + WG) {
+    acc = acc + f32(a[k]) * dequant_weight(b_row_base + k);
+  }
+  partials[tid] = acc;
+  // Ceil-halving tree reduce: WG is device-bounded and need not be
+  // a power of two, so the live count is carried explicitly. Plain `stride /= 2`
+  // would pair a live lane with an already-consumed slot at the first odd count.
+  var count = WG;
+  loop {
+    workgroupBarrier();
+    let half = (count + 1u) / 2u;
+    if (tid + half < count) {
+      partials[tid] = partials[tid] + partials[tid + half];
+    }
+    if (count == 1u) {
+      break;
+    }
+    count = half;
+  }
+  if (tid == 0u) {
+    y[col] = {{ aScalar }}(partials[0]);
+  }
+}

build/webgpu/matmul-bnb4-sgmat.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,128 @@

+// Large aligned prefill tier. The packed B tile is dequantized once per
+// workgroup, then reused by subgroup-matrix f32 multiply-accumulate operations.
+// This specialization operates on full A, B, and output tiles, so both
+// subgroupMatrixLoad and the direct output stores are in bounds.
+enable subgroups;
+enable chromium_experimental_subgroup_matrix;
+{% if usesF16 %}
+// The widened path keeps the activations and the output in f32 for the matrix
+// units, but the quantization scales stay in the request's own dtype.
+enable f16;
+{% endif %}
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const BLOCK_SIZE: u32 = {{ blockSize }}u;
+const TILE_ROWS: u32 = {{ tileRows }}u;
+const TILE_COLS: u32 = {{ tileCols }}u;
+const TILE_K: u32 = {{ tileK }}u;
+const SUB_ROWS: u32 = {{ subRows }}u;
+const SUB_COLS: u32 = {{ subCols }}u;
+const LOAD_WIDTH: u32 = {{ loadWidth }}u;
+var<workgroup> tile_B: array<f32, {{ tileCols * tileK }}>;
+/* bitsandbytes 4-bit weight decoding.
+   Only the selected FP4 or NF4 codebook is emitted. The codebook uses
+   var<private> rather than const because its runtime index requires a memory
+   location in WGSL.
+   FP4 is sign-magnitude (bit 3 is sign), so its table is non-monotonic. NF4 is
+   the ascending normal-float quantile codebook. Each stored byte is carried in
+   one widened u32. A flat weight index maps to byte flat >> 1, and an even
+   index reads the high nibble, matching (vi0 << 4) | vi1 packing. */
+{% if quantType == 0 %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  0.0, 5.208333333e-03, 0.66666667, 1.0,
+  0.33333333, 0.5, 0.16666667, 0.25,
+  -0.0, -5.208333333e-03, -0.66666667, -1.0,
+  -0.33333333, -0.5, -0.16666667, -0.25
+);
+{%- else %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
+  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
+  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
+  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
+);
+{%- endif %}
+fn load_b(tile_base: u32, k_base: u32, row: u32, col_group: u32) {
+  let col = col_group * LOAD_WIDTH;
+  let flat = (tile_base + row) * K + k_base + col;
+  let scale = f32(absmax[flat / BLOCK_SIZE]);
+  let tile_offset = row * TILE_K + col;
+  // Eligibility guarantees that each aligned LOAD_WIDTH slice stays in one
+  // quantization block and begins on an even weight. The loader widens each
+  // packed byte once and reuses both nibbles.
+  for (var i = 0u; i < LOAD_WIDTH; i = i + 2u) {
+    let packed = b[(flat + i) >> 1u];
+    let high = (packed >> 4u) & 0xFu;
+    let low = packed & 0xFu;
+    tile_B[tile_offset + i] = quant_map[high] * scale;
+    tile_B[tile_offset + i + 1u] = quant_map[low] * scale;
+  }
+}
+@compute @workgroup_size({{ workgroupSize }}, 1, 1)
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_index) local_idx: u32,
+  @builtin(subgroup_size) subgroup_size: u32
+) {
+  let a_base = workgroup_id.y * TILE_ROWS;
+  let b_base = workgroup_id.x * TILE_COLS;
+  let subgroup_id = local_idx / subgroup_size;
+  let subtile_idx = subgroup_id / {{ rowSubtiles }}u;
+  let subtile_idy = subgroup_id % {{ rowSubtiles }}u;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+{% for row in range(rowMatrices) %}{% for col in range(colMatrices) %}
+  var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, {{ matrixSize }}, {{ matrixSize }}>;
+{% endfor %}{% endfor %}
+  for (var k_base = 0u; k_base < K; k_base += TILE_K) {
+    load_b(b_base, k_base, local_idx / {{ bLoadsPerRow }}u, local_idx % {{ bLoadsPerRow }}u);
+    workgroupBarrier();
+    for (var step = 0u; step < TILE_K; step += {{ matrixSize }}u) {
+      let matrix_a_offset = (a_base + base_A) * K + k_base + step;
+{% for row in range(rowMatrices) %}
+      let matA{{ row }} = subgroupMatrixLoad<subgroup_matrix_left<f32, {{ matrixSize }}, {{ matrixSize }}>>(
+        &a, matrix_a_offset + {{ row * matrixSize }}u * K, false, K
+      );
+{% endfor %}
+      let matrix_b_offset = subtile_idx * SUB_COLS * TILE_K + step;
+{% for col in range(colMatrices) %}
+      let matB{{ col }} = subgroupMatrixLoad<subgroup_matrix_right<f32, {{ matrixSize }}, {{ matrixSize }}>>(
+        &tile_B, matrix_b_offset + {{ col * matrixSize }}u * TILE_K, true, TILE_K
+      );
+{% endfor %}
+{% for row in range(rowMatrices) %}{% for col in range(colMatrices) %}
+      matC{{ row }}{{ col }} = subgroupMatrixMultiplyAccumulate(matA{{ row }}, matB{{ col }}, matC{{ row }}{{ col }});
+{% endfor %}{% endfor %}
+    }
+    workgroupBarrier();
+  }
+  let output_offset = (a_base + base_A) * N + b_base + base_B;
+{% for row in range(rowMatrices) %}{% for col in range(colMatrices) %}
+  subgroupMatrixStore(
+    &y,
+    output_offset + {{ row * matrixSize }}u * N + {{ col * matrixSize }}u,
+    matC{{ row }}{{ col }},
+    false,
+    N
+  );
+{% endfor %}{% endfor %}
+}

build/webgpu/matmul-bnb4-tiled.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,168 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+// Prefill: register-blocked 64x64 shared-memory tiled GEMM (16x8 threads, each
+// holding an 8x4 micro-tile, with K stepped in BK chunks).
+//
+// Each weight is dequantized once into the shared B tile and reused across the
+// row block. Both operands are stored K-major — A is [M, K] and the packed
+// weight is [N, K] — so both tiles group four K values per vector word, every
+// staging read is contiguous, and one micro-tile step reads TM + TN vector
+// words to issue TM * TN * 4 fused multiply-adds through dot().
+const BK: u32 = {{ tileK }}u;
+const BM: u32 = 64u;
+const BN: u32 = 64u;
+const TM: u32 = 8u;
+const TN: u32 = 4u;
+const WG_X: u32 = 16u;
+const WG_SIZE: u32 = 128u;
+const K_VECS: u32 = BK / 4u;
+// Stage tileA at its input storage width; widening it on shared load is
+// bit-identical. tileB remains f32 to avoid narrowing dequantized NF4/FP4
+// values before accumulation.
+{% set aTile = "f16" if usesF16 else "f32" %}
+var<workgroup> tileA: array<array<vec4<{{ aTile }}>, K_VECS>, BM>;
+var<workgroup> tileB: array<array<vec4<f32>, K_VECS>, BN>;
+/* bitsandbytes 4-bit weight decoding.
+   Only the selected FP4 or NF4 codebook is emitted. The codebook uses
+   var<private> rather than const because its runtime index requires a memory
+   location in WGSL.
+   FP4 is sign-magnitude (bit 3 is sign), so its table is non-monotonic. NF4 is
+   the ascending normal-float quantile codebook. Each stored byte is carried in
+   one widened u32. A flat weight index maps to byte flat >> 1, and an even
+   index reads the high nibble, matching (vi0 << 4) | vi1 packing. */
+{% if quantType == 0 %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  0.0, 5.208333333e-03, 0.66666667, 1.0,
+  0.33333333, 0.5, 0.16666667, 0.25,
+  -0.0, -5.208333333e-03, -0.66666667, -1.0,
+  -0.33333333, -0.5, -0.16666667, -0.25
+);
+{%- else %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
+  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
+  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
+  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
+);
+{%- endif %}
+// Four consecutive dequantized weights starting at flat index `flat` into the
+// [N, K] weight. Consecutive indices share a stored byte, so the run costs two
+// or three byte loads instead of four; the parity of `flat` decides which
+// nibble opens the run.
+fn dequant_weight4(flat: u32) -> vec4<f32> {
+  let base = flat >> 1u;
+  let byte0 = b[base];
+  let byte1 = b[base + 1u];
+  let byte2 = b[base + 2u];
+  var codes: vec4<u32>;
+  if ((flat & 1u) == 0u) {
+    codes = vec4<u32>((byte0 >> 4u) & 0xFu, byte0 & 0xFu, (byte1 >> 4u) & 0xFu, byte1 & 0xFu);
+  } else {
+    codes = vec4<u32>(byte0 & 0xFu, (byte1 >> 4u) & 0xFu, byte1 & 0xFu, (byte2 >> 4u) & 0xFu);
+  }
+  // block_size is a power of two of at least 16, so a four-weight run stays
+  // inside one block unless it starts within three of the block's end.
+  var scales: vec4<f32>;
+  if (flat % params.blockSize + 3u < params.blockSize) {
+    scales = vec4<f32>(f32(absmax[flat / params.blockSize]));
+  } else {
+    scales = vec4<f32>(
+      f32(absmax[flat / params.blockSize]),
+      f32(absmax[(flat + 1u) / params.blockSize]),
+      f32(absmax[(flat + 2u) / params.blockSize]),
+      f32(absmax[(flat + 3u) / params.blockSize])
+    );
+  }
+  return vec4<f32>(quant_map[codes.x], quant_map[codes.y], quant_map[codes.z], quant_map[codes.w]) * scales;
+}
+@compute @workgroup_size(16, 8, 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 * WG_X + lid.x;
+  let aRow = lid.y * TM;
+  let bCol = lid.x * TN;
+{% for row in range(8) %}
+  var acc{{ row }} = vec4<f32>(0.0);
+{% endfor %}
+  let numTiles = (params.K + BK - 1u) / BK;
+  for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
+    let kBase = kt * BK;
+    // Each lane stages whole vector words: four consecutive K values of one A
+    // row and of one dequantized weight row.
+    for (var linear = li; linear < BM * K_VECS; linear += WG_SIZE) {
+      let ar = linear / K_VECS;
+      let ak4 = linear % K_VECS;
+      let am = mBase + ar;
+      let ak = kBase + ak4 * 4u;
+      var av = vec4<{{ aTile }}>({{ aTile }}(0.0));
+      if (am < params.rows) {
+        let aFlat = am * params.K + ak;
+        if (ak + 3u < params.K) {
+          av = vec4<{{ aTile }}>({{ aTile }}(a[aFlat]), {{ aTile }}(a[aFlat + 1u]), {{ aTile }}(a[aFlat + 2u]), {{ aTile }}(a[aFlat + 3u]));
+        } else {
+          for (var t = 0u; t < 4u; t = t + 1u) {
+            if (ak + t < params.K) {
+              av[t] = {{ aTile }}(a[aFlat + t]);
+            }
+          }
+        }
+      }
+      tileA[ar][ak4] = av;
+    }
+    for (var linear = li; linear < BN * K_VECS; linear += WG_SIZE) {
+      let br = linear / K_VECS;
+      let bk4 = linear % K_VECS;
+      let bn = nBase + br;
+      let bk = kBase + bk4 * 4u;
+      var bv = vec4<f32>(0.0);
+      if (bn < params.N) {
+        bv = dequant_weight4(bn * params.K + bk);
+        if (bk + 3u >= params.K) {
+          for (var t = 0u; t < 4u; t = t + 1u) {
+            if (bk + t >= params.K) {
+              bv[t] = 0.0;
+            }
+          }
+        }
+      }
+      tileB[br][bk4] = bv;
+    }
+    workgroupBarrier();
+    for (var kv: u32 = 0u; kv < K_VECS; kv = kv + 1u) {
+{% for column in range(4) %}
+      let bv{{ column }} = tileB[bCol + {{ column }}u][kv];
+{% endfor %}
+{% for row in range(8) %}
+      let av{{ row }} = vec4<f32>(tileA[aRow + {{ row }}u][kv]);
+      acc{{ row }} += vec4<f32>({% for column in range(4) %}dot(av{{ row }}, bv{{ column }}){% if not loop.last %}, {% endif %}{% endfor %});
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% for row in range(8) %}
+{% for column in range(4) %}
+  {
+    let m = mBase + aRow + {{ row }}u;
+    let n = nBase + bCol + {{ column }}u;
+    if (m < params.rows && n < params.N) { y[m * params.N + n] = {{ aScalar }}(acc{{ row }}.{{ ["x", "y", "z", "w"][column] }}); }
+  }
+{% endfor %}
+{% endfor %}
+}

build/webgpu/matmul-bnb4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,65 @@

+{% if usesF16 %}
+enable f16;
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ workgroupSize }}u;
+/* bitsandbytes 4-bit weight decoding.
+   Only the selected FP4 or NF4 codebook is emitted. The codebook uses
+   var<private> rather than const because its runtime index requires a memory
+   location in WGSL.
+   FP4 is sign-magnitude (bit 3 is sign), so its table is non-monotonic. NF4 is
+   the ascending normal-float quantile codebook. Each stored byte is carried in
+   one widened u32. A flat weight index maps to byte flat >> 1, and an even
+   index reads the high nibble, matching (vi0 << 4) | vi1 packing. */
+{% if quantType == 0 %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  0.0, 5.208333333e-03, 0.66666667, 1.0,
+  0.33333333, 0.5, 0.16666667, 0.25,
+  -0.0, -5.208333333e-03, -0.66666667, -1.0,
+  -0.33333333, -0.5, -0.16666667, -0.25
+);
+{%- else %}
+var<private> quant_map: array<f32, 16> = array<f32, 16>(
+  -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
+  -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
+  0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224,
+  0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0
+);
+{%- endif %}
+// One dequantized weight at flat index `flat` into the [N, K] weight.
+fn dequant_weight(flat: u32) -> f32 {
+  let byte = b[flat >> 1u];
+  let code = select(byte & 0xFu, (byte >> 4u) & 0xFu, (flat & 1u) == 0u);
+  return quant_map[code] * f32(absmax[flat / params.blockSize]);
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat output-element index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit. When nwg.y == 1 this
+  // reduces to gid.x; the index >= total guard drops the tail.
+  let index = gid.x + gid.y * nwg.x * WG;
+  let total = params.rows * params.N;
+  if (index >= total) {
+    return;
+  }
+  let row = index / params.N;
+  let col = index % params.N;
+  var acc = 0.0;
+  let b_row_base = col * params.K;
+  for (var k: u32 = 0u; k < params.K; k = k + 1u) {
+    acc = acc + f32(a[row * params.K + k]) * dequant_weight(b_row_base + k);
+  }
+  y[index] = {{ aScalar }}(acc);
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,22 @@

+{
+  "name": "com.microsoft.MatMulBnb4",
+  "id": "_com_microsoft_matmulbnb4_webgpu_0097122",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "doteNv/ofPVMbmfcGye+YHqSiBcZ65LVeDl4oDWa1Ec=",
+      "cast-scalar-x4.wgsl.jinja": "uplwFdE8dJSffJop4txQHZgeGDDJ5K9+F3V4n1CotVo=",
+      "manifest.json": "c2kZSTr7EnI+nB6vtDMF30I5Io1qYvCJwr69mQJxeD8=",
+      "matmul-bnb4-gemv.wgsl.jinja": "7e4ZY+raLAMG/wmXm72XpSzPsR6adHvvrcKxv2ZKj/4=",
+      "matmul-bnb4-sgmat.wgsl.jinja": "tm3SCGjsEM1jwLUR0W4WdjU75y8NcOSsa/1vt81YryE=",
+      "matmul-bnb4-tiled.wgsl.jinja": "+pb1pUgFU+RLITtJ+zrnTDl3LUmU7sZDGqIdyKtAmn4=",
+      "matmul-bnb4.wgsl.jinja": "xzeJm4jFqbXih7plcQup4jUQsRw4h8HpLzcAxWR2aic=",
+      "test.json": "1oYDh6veLTUP+o7nE+G7nU1FMK88nLKArXE/FcbOfTo="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MatMulBnb4" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,434 @@

+{
+  "op": "com.microsoft.MatMulBnb4",
+  "fixtureArrays": {
+    "bnb4_weight_cycle_b_t": [1, 35, 69, 103, 137, 171, 205, 239, 254, 220, 186, 152, 118, 84, 50, 16, 90]
+  },
+  "cases": [
+    {
+      "name": "nf4_scalar_small",
+      "attrs": { "K": 16, "N": 4, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "fp4_scalar_small",
+      "attrs": { "K": 16, "N": 4, "block_size": 16, "quant_type": 0 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [32],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [3, 4], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_scalar_block_straddle",
+      "attrs": { "K": 24, "N": 3, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 24],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [36],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [2, 3], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_scalar_odd_k",
+      "attrs": { "K": 33, "N": 5, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [3, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [83],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [11],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [3, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_gemv_decode",
+      "attrs": { "K": 128, "N": 64, "block_size": 64, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [4096],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "fp4_gemv_decode",
+      "attrs": { "K": 64, "N": 32, "block_size": 32, "quant_type": 0 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1024],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 32], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_gemv_ragged_k",
+      "attrs": { "K": 100, "N": 12, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 100],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [600],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [75],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [1, 12], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_tiled_prefill",
+      "attrs": { "K": 64, "N": 48, "block_size": 32, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1536],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [32, 48], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "fp4_tiled_prefill",
+      "attrs": { "K": 64, "N": 48, "block_size": 32, "quant_type": 0 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [32, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [1536],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [96],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [32, 48], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "nf4_sgmat_tall_aligned",
+      "attrs": { "K": 64, "N": 64, "block_size": 32, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [2048],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [128, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "fp4_sgmat_aligned",
+      "attrs": { "K": 64, "N": 64, "block_size": 32, "quant_type": 0 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [64, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [2048],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [64, 64], "tolerance": 0.00002 } }
+    },
+    {
+      "name": "nf4_tiled_tail",
+      "attrs": { "K": 33, "N": 19, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [20, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [314],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [40],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [20, 19], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_tiled_tail_bk32",
+      "attrs": { "K": 33, "N": 19, "block_size": 16, "quant_type": 1 },
+      "tunables": { "PORTABLE_TILE_K": 32 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [20, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [314],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [40],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [20, 19], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "fp4_tiled_tail_bk64",
+      "attrs": { "K": 33, "N": 19, "block_size": 16, "quant_type": 0 },
+      "tunables": { "PORTABLE_TILE_K": 64 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [20, 33],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [314],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float32",
+          "shape": [40],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float32", "shape": [20, 19], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "nf4_f16_scalar",
+      "attrs": { "K": 32, "N": 8, "block_size": 16, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [3, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [128],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float16",
+          "shape": [16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [3, 8], "tolerance": 0.02 } }
+    },
+    {
+      "name": "nf4_f16_gemv",
+      "attrs": { "K": 64, "N": 16, "block_size": 32, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [512],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float16",
+          "shape": [32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.02 } }
+    },
+    {
+      "name": "nf4_f16_tiled",
+      "attrs": { "K": 32, "N": 16, "block_size": 32, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [16, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [256],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float16",
+          "shape": [16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [16, 16], "tolerance": 0.02 } }
+    },
+    {
+      "name": "nf4_f16_tiled_longk_m32_k512_n32",
+      "attrs": { "K": 512, "N": 32, "block_size": 64, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [32, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [8192],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float16",
+          "shape": [256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [32, 32], "tolerance": 0.001, "relTolerance": 0.001 } },
+      "provenance": {
+        "notes": "The other f16 fixtures all run K <= 64, where an f16 accumulator is indistinguishable from an f32 one. This one accumulates over K=512 so the tiled path's f32 accumulation is actually pinned."
+      }
+    },
+    {
+      "name": "nf4_sgmat_widened_f16",
+      "provenance": {
+        "notes": "A float16 request on the aligned subgroup-matrix shape. The matrix units take float32 operands and accumulate in float32, so this path widens the activations either side of the multiply rather than using the device's float16 configuration, whose result component is also float16. It is the only case that reaches the widening passes."
+      },
+      "attrs": { "K": 64, "N": 64, "block_size": 32, "quant_type": 1 },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "bT": {
+          "dtype": "uint8",
+          "shape": [2048],
+          "data": { "kind": "cycle", "values": { "$ref": "#/fixtureArrays/bnb4_weight_cycle_b_t" } }
+        },
+        "absmaxT": {
+          "dtype": "float16",
+          "shape": [128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.07, "scale": 0.02, "offset": 0.06 }
+        }
+      },
+      "outputs": { "yT": { "dtype": "float16", "shape": [128, 64], "tolerance": 0.002, "relTolerance": 0.002 } }
+    }
+  ]
+}