MatMulNBits (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::MatMulNBits

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1, until opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeMatMulNBits()

Peak memory

onnx_light_cpu::ComputePeakMemoryMatMulNBits()

Fusion patterns

onnx_light_cpu::MatMulNBitsBiasFusionPattern

Gradient

onnx_light_cpu::RegisterCustomOperatorGradients()

LightOpSchema#

Version: 1

Restricted compatibility subset of ONNX Runtime MatMulNBits. The CPU implementation accepts matching FLOAT, FLOAT16, or BFLOAT16 activations, scales, optional bias, and output; UINT8 packed INT2, INT4, or INT8 weights; power-of-two block_size in [16, UINT32_MAX], weight_prepacked=0, and accuracy_level 0 or 4. B may instead be an ORT MatMulNBits EncodedValueProto (INT2/4/8) or an onnx-light MATMULNBITS_INT3 EncodedValueProto, in which case scales and zero_points must be omitted and the encoded packing is consumed without provider repacking. INT3 requires block_size=32. This encoded form is inference-only. zero_points may be packed UINT8 or match A’s floating type; g_idx must not be supplied. Packed B has shape (N, ceil(K / block_size), block_size * bits / 8), scales has shape (N, ceil(K / block_size)) or its flattened equivalent, and bias has shape (N).

Inputs

  • A (T1): FLOAT, FLOAT16, or BFLOAT16 tensor of rank at least one ending in K.

  • B (T2): UINT8 packed weights with shape (N, ceil(K / block_size), block_size * bits / 8), or an ORT MatMulNBits EncodedValueProto (required for INT3).

  • scales (T1): Block scales with shape (N, ceil(K / block_size)) or its flattened equivalent; omitted when B is encoded.

  • zero_points (T3): Optional packed UINT8 zero points, byte-aligned per column, or floating zero points with scales’ shape and dtype. Omit when B is encoded.

  • g_idx (T4): Unsupported deprecated upstream group-index slot; must be omitted.

  • bias (T1): Optional bias with shape (N) and A’s type.

Outputs

  • Y (T1): Output with A’s type, leading dimensions, and final dimension N.

Attributes

  • K (INT): Positive input feature count. (required)

  • N (INT): Positive output feature count. (required)

  • accuracy_level (INT): Accuracy level; 0 and 4 are supported. (optional, default: 0)

  • bits (INT): Quantized weight bit width; 2, 4, 8 or encoded-only 3. (optional, default: 4)

  • block_size (INT): Power of two in [16, UINT32_MAX]; INT3 requires 32. (required)

  • weight_prepacked (INT): Prepacked-weight flag; only 0 is supported. (optional, default: 0)

Type constraints

  • T1: Constrain A, scales, bias, and Y to matching non-double floating-point tensors. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • T2: Constrain packed B to UINT8 or a structured encoded value. Allowed types: tensor(struct), tensor(uint8).

  • T3: Constrain zero_points to packed UINT8 or A’s floating-point type. Allowed types: tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8).

  • T4: Retain the upstream g_idx type slot; this subset rejects a supplied value. Allowed types: tensor(int32).