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 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 — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
cast-scalar-x4.wgsl.jinja
matmul-bnb4-gemv.wgsl.jinja
matmul-bnb4-sgmat.wgsl.jinja
matmul-bnb4-tiled.wgsl.jinja
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.

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,
  },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.