Tutorial: MatMulNBits Custom Op Replacement

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/custom_op/MatMulNBits_custom_op_tutorial.ipynb.

MatMulNBits Custom Op Replacement

Replace ONNX MatMulNBits (4-bit quantized matmul with per-block scales) with the Quadric custom op linalg::matrixMulMatrix, which computes Y = A @ B_dequantized.

This tutorial walks through loading a MatMulNBits ONNX model, computing tensor ranges from dequantized weights and ORT output ranges, replacing each MatMulNBits node with a Quadric custom op using replace_subgraph_by_edges, and finally compiling with CGC and validating against ORT.

Model: MatMulNBits single-layer (4-bit quantized, per-block scales)

Notebook Outline

• 1. Setup
• 2. Model Preparation
• 3. Custom Op Replacement
• 4. Compile with Chimera Graph Compiler

1. Setup

Configuration

MODEL_PATH = "./matmulnbits_single_layer.onnx"
OUTPUT_PATH = MODEL_PATH[:-5] + "_replaced.onnx"
TRANGES_PATH = MODEL_PATH[:-5] + "_replaced.tranges"

Imports

import onnx
from onnx import numpy_helper, version_converter
import numpy as np
from tvm.contrib.epu.chipy.param_utils import max_fraction_bits
import tvm.contrib.epu.graphutils as graphutils
from tvm.contrib.epu.graphutils import CustomOpReplacer
from functools import partial

Helper Functions

def get_tensor_data(tensor_name, model):
    """Get tensor data from model initializers."""
    for init in model.graph.initializer:
        if init.name == tensor_name:
            return numpy_helper.to_array(init)
    return None


def find_matmulnbits_nodes(model):
    """Find all MatMulNBits nodes and return their info dicts."""
    matmulnbits_nodes = []
    for node in model.graph.node:
        if node.op_type == "MatMulNBits" and node.domain == "com.microsoft":
            attrs = {}
            for attr in node.attribute:
                if attr.name in ["K", "N", "bits", "block_size"]:
                    attrs[attr.name] = attr.i
            matmulnbits_nodes.append(
                {
                    "node": node,
                    "input_a": node.input[0],
                    "input_b": node.input[1],
                    "input_scales": node.input[2],
                    "input_zero_points": node.input[3] if len(node.input) > 3 else None,
                    "output": node.output[0],
                    "attributes": attrs,
                    "name": node.name,
                }
            )
    return matmulnbits_nodes


def unpack_v8i4_weights(packed_weights):
    """Unpack V8I4 format: two 4-bit values per byte (low nibble first)."""
    flat_packed = packed_weights.flatten()
    unpacked = np.zeros(len(flat_packed) * 2, dtype=np.uint8)
    unpacked[0::2] = flat_packed & 0x0F
    unpacked[1::2] = (flat_packed >> 4) & 0x0F
    return unpacked


def preprocess_weights(weights_data, K, N, block_size, scales_data=None):
    """Preprocess MatMulNBits weights: unpack, reshape, sign-convert, apply per-block scales.

    ONNX stores unsigned [0,15]; we subtract 8 to get signed [-8,7] (equivalent to
    the signed int4 in the GPNPU sweep tests). Then dequantize: weight_float = int4 * scale[block, n].
    """
    n_blocks = K // block_size

    # Unpack V8I4 [N, n_blocks, block_size/2] and reshape to [K, N]
    int4_weights = unpack_v8i4_weights(weights_data)
    int4_weights_2d = int4_weights[: K * N].reshape(N, K).T

    # Convert unsigned [0,15] to signed [-8,7]
    int4_signed = int4_weights_2d.astype(np.int8) - 8

    if scales_data is not None:
        # ONNX scales layout: [N, n_blocks]  (sweep_runner uses [n_blocks, N] — transposed)
        scales_2d = scales_data.reshape(N, n_blocks)
        dequant = np.zeros((K, N), dtype=np.float32)
        for b in range(n_blocks):
            k_start = b * block_size
            k_end = k_start + block_size
            dequant[k_start:k_end, :] = (
                int4_signed[k_start:k_end, :].astype(np.float32) * scales_2d[:, b]
            )
        return dequant
    else:
        return int4_signed.astype(np.float32)


def make_process_matmulnbits_consts(K, N, block_size, bits, scales_data=None):
    """Create process_const_callback that dequantizes weights during replacement."""

    def process_matmulnbits_consts(param_dict, const_inputs, idx):
        reversed_param_dict = {}
        for param_name, tensor_name in param_dict.items():
            reversed_param_dict.setdefault(tensor_name, []).append(param_name)

        results_list = []
        for _key, const_data in const_inputs.items():
            tensor_name = const_data[0]
            tensor_data = const_data[1]
            if tensor_name in reversed_param_dict:
                param_name = reversed_param_dict[tensor_name][0]
                if "weights" in param_name:
                    processed = preprocess_weights(
                        tensor_data, K, N, block_size, scales_data=scales_data
                    )
                    results_list.append((f"{param_name}_{idx}", processed))
                else:
                    results_list.append((f"{param_name}_{idx}", tensor_data))
        return {idx: (name, data) for idx, (name, data) in enumerate(results_list)}

    return process_matmulnbits_consts


def compute_error_metrics(reference, prediction):
    """Compute RMSE, correlation, and max absolute error between two tensors."""
    ref_flat = reference.flatten()
    pred_flat = prediction.flatten()
    diff = pred_flat - ref_flat
    return {
        "rmse": float(np.sqrt(np.mean(diff**2))),
        "correlation": float(np.corrcoef(ref_flat, pred_flat)[0, 1]),
        "max_abs_error": float(np.abs(diff).max()),
    }


def replace_matmulnbits_with_custom_op(model_path, tranges_dict):
    """Replace all MatMulNBits nodes in an ONNX model with Quadric custom ops.

    Args:
        model_path: Path to the ONNX model file.
        tranges_dict: Dict of tensor name -> [min, max] ranges. Weight ranges will be
            added automatically. Output ranges are looked up by name; if missing, frac
            bits are estimated from the dequantized weight range.

    Returns:
        Tuple of (replaced_model_path, tranges_path, original_model, custom_model).
    """
    model = onnx.load(model_path)
    nodes = find_matmulnbits_nodes(model)
    print(f"Found {len(nodes)} MatMulNBits node(s):")
    for i, info in enumerate(nodes):
        attrs = info["attributes"]
        print(
            f"  Node {i}: {info['name']} - K={attrs['K']}, N={attrs['N']}, "
            f"block_size={attrs['block_size']}, bits={attrs['bits']}"
        )

    # Add dequantized weight ranges to tranges
    output_frac_bits_list = []
    for idx, info in enumerate(nodes):
        K = info["attributes"]["K"]
        N = info["attributes"]["N"]
        block_size = info["attributes"]["block_size"]

        weights_data = get_tensor_data(info["input_b"], model)
        scales_data = get_tensor_data(info["input_scales"], model)
        dequant_weights = preprocess_weights(weights_data, K, N, block_size, scales_data)
        max_abs_w = max(abs(float(dequant_weights.min())), abs(float(dequant_weights.max())))
        tranges_dict[f"weights_{idx + 1}"] = [-max_abs_w, max_abs_w]
        print(
            f"  weights_{idx + 1} range: [{dequant_weights.min():.4f}, {dequant_weights.max():.4f}]"
        )

        # Compute output frac bits from tranges (use original model output names)
        output_name = info["output"]
        if output_name in tranges_dict:
            output_range = tranges_dict[output_name]
            max_abs_out = max(abs(output_range[0]), abs(output_range[1]))
        else:
            # Estimate from weight range * assumed input range
            max_abs_out = max_abs_w * K * 0.5
            print(
                f"    Warning: output '{output_name}' not in tranges, estimating frac_bits from weight range"
            )
        output_frac_bits_list.append(max_fraction_bits(max_abs_out))
        print(f"    output_frac_bits={output_frac_bits_list[-1]}")

    # Save tranges
    tranges_path = model_path[:-5] + "_replaced.tranges"
    with open(tranges_path, "w") as f:
        json.dump(tranges_dict, f, indent=2)

    # Convert to opset 16 and replace nodes
    model_opset16 = version_converter.convert_version(model, 16)
    nodes_opset16 = find_matmulnbits_nodes(model_opset16)

    util = CustomOpReplacer(model_opset16)
    custom_model = model_opset16

    for idx, info in enumerate(nodes_opset16):
        K = info["attributes"]["K"]
        N = info["attributes"]["N"]
        block_size = info["attributes"]["block_size"]
        bits = info["attributes"]["bits"]

        scales_data = get_tensor_data(info["input_scales"], model_opset16)
        if scales_data is None:
            scales_data = get_tensor_data(info["input_scales"], model)

        param_dict = {"weights": info["input_b"]}
        out_edges = [info["output"]]
        in_edges = [info["input_a"]]

        const_callback = make_process_matmulnbits_consts(
            K, N, block_size, bits, scales_data=scales_data
        )
        ccl_function = "linalg::matrixMulMatrix"

        _, custom_model = util.replace_subgraph_by_edges(
            out_edges,
            in_edges,
            ccl_function,
            element_wise=False,
            fixed_point_frac_bits=[output_frac_bits_list[idx]],
            keep_constants=list(param_dict.values()),
            process_const_callback=partial(const_callback, param_dict),
            allow_io_in_l2_mem=False,
            needs_iter_var=False,
        )
        print(
            f"Replaced node {idx}: {info['name']} -> {ccl_function} (frac_bits={output_frac_bits_list[idx]})"
        )

    # Save replaced model
    output_path = model_path[:-5] + "_replaced.onnx"
    onnx.save(custom_model, output_path)
    print(f"\nSaved replaced model to {output_path}")

    return output_path, tranges_path, model, custom_model

2. Model Preparation

Load the ONNX model, compute tranges, and identify all MatMulNBits nodes.

model = onnx.load(MODEL_PATH)
nodes = find_matmulnbits_nodes(model)

for i, node_info in enumerate(nodes):
    attrs = node_info["attributes"]
    print(
        f"Node {i}: K={attrs['K']}, N={attrs['N']}, block_size={attrs['block_size']}, bits={attrs['bits']}"
    )

K = nodes[0]["attributes"]["K"]

Node 0: K=384, N=1152, block_size=128, bits=4

3. Custom Op Replacement

Compute tensor ranges and replace MatMulNBits nodes with linalg::matrixMulMatrix.

import json
import onnxruntime as ort

node_info = nodes[0]
K = node_info["attributes"]["K"]
N = node_info["attributes"]["N"]
block_size = node_info["attributes"]["block_size"]

weights_data = get_tensor_data(node_info["input_b"], model)
scales_data = get_tensor_data(node_info["input_scales"], model)

## Dequantized weight range
dequant_weights = preprocess_weights(weights_data, K, N, block_size, scales_data)
max_abs_weight = max(abs(float(dequant_weights.min())), abs(float(dequant_weights.max())))

## ORT output range from random inputs
sess = ort.InferenceSession(MODEL_PATH)
max_abs_y = 0
for _ in range(10):
    dummy_a = np.random.uniform(-10, 10, (1, 1, K)).astype(np.float32)
    ort_out = sess.run(None, {"A": dummy_a})[0]
    max_abs_y = max(max_abs_y, abs(float(ort_out.min())), abs(float(ort_out.max())))

tranges = {
    "A": [-10.0, 10.0],
    "Y": [-max_abs_y * 1.5, max_abs_y * 1.5],
}

print(f"Weight range: [{dequant_weights.min():.4f}, {dequant_weights.max():.4f}]")
print(f"Output Y range estimate: [{-max_abs_y * 1.5:.4f}, {max_abs_y * 1.5:.4f}]")

OUTPUT_PATH, TRANGES_PATH, _, custom_model = replace_matmulnbits_with_custom_op(MODEL_PATH, tranges)

print(f"\nModel nodes after replacement:")
for node in custom_model.graph.node:
    print(f"  {node.op_type} ({node.name}): inputs={list(node.input)}, outputs={list(node.output)}")

Weight range: [-27.3287, 27.6624]
Output Y range estimate: [-4902.7972, 4902.7972]
Found 1 MatMulNBits node(s):
  Node 0: matmul_nbits_layer - K=384, N=1152, block_size=128, bits=4
  weights_1 range: [-27.3287, 27.6624]
    output_frac_bits=18
Replaced node 0: matmul_nbits_layer -> linalg::matrixMulMatrix (frac_bits=18)

Saved replaced model to ./matmulnbits_single_layer_replaced.onnx

Model nodes after replacement:
  QuadricCustomOp (CustomOp/linalg::matrixMulMatrix0): inputs=['A', 'weights_1'], outputs=['Y']

4. Compile with Chimera Graph Compiler

Create a CGC job for the replaced model and validate ISS output against ORT.

import matplotlib

matplotlib.use("Agg")

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from tvm.contrib.epu.chimera_job.hw_config import HWConfig

## Add output 'Y' to value_info for ORT validation compatibility
original_model = onnx.load(MODEL_PATH)
output_Y = original_model.graph.output[0]
y_in_value_info = any(vi.name == "Y" for vi in original_model.graph.value_info)
if not y_in_value_info:
    original_model.graph.value_info.append(output_Y)
    ORT_MODEL_PATH = MODEL_PATH[:-5] + "_ort_fixed.onnx"
    onnx.save(original_model, ORT_MODEL_PATH)
else:
    ORT_MODEL_PATH = MODEL_PATH

hw_config = HWConfig(product="QC-N")
cgc_job = ChimeraJob(
    model_p=OUTPUT_PATH,
    hw_config=hw_config,
    bypass_quantization_checks=True,
    validate_iss=True,
    trange_file=TRANGES_PATH,
    onnx_ort_override_p=ORT_MODEL_PATH,
)

inputs = {"A": np.random.uniform(-10.0, 10.0, size=(1, 1, K)).astype(np.float32)}

cgc_job.compile()
print(cgc_job)


## Workaround: run ORT and ISS separately and compare outputs directly.
ort_outputs = cgc_job.run_onnx_inf_session(inputs=inputs)
iss_outputs = cgc_job.run_inference_harness(inputs=inputs, compare_ort=False)

ort_Y = ort_outputs["Y"]
iss_Y = iss_outputs["Y"].tensor.reshape(ort_Y.shape)

metrics = compute_error_metrics(ort_Y, iss_Y)
print(f"\n=== ORT vs ISS Comparison ===")
print(f"ORT output range: [{ort_Y.min():.4f}, {ort_Y.max():.4f}]")
print(f"ISS output range: [{iss_Y.min():.4f}, {iss_Y.max():.4f}]")
print(f"RMSE:          {metrics['rmse']:.4f}")
print(f"Correlation:   {metrics['correlation']:.6f}")
print(f"Max abs error: {metrics['max_abs_error']:.4f}")

assert (
    metrics["rmse"] < 1e-3
), "MatMul Error : Error is too high, please check malmutnbits implementation"

2026-06-19 03:53 - INFO - epu - chimera_job - Overriding onnx model ort execution with the onnx model: ./matmulnbits_single_layer_ort_fixed.onnx
2026-06-19 03:53 - INFO - epu - chimera_job - START==================================onnx_ingest
2026-06-19 03:53 - INFO - epu - chimera_job - Numerical ranges provided
2026-06-19 03:53 - INFO - epu - codegen - START===============================optimize_relay
2026-06-19 03:53 - INFO - epu - codegen - START====================quantize_to_cpu_runnable_fx
2026-06-19 03:53 - INFO - epu - fx - 

Source name                        Op                             Output 0 Range         Output 0 Frac Bits
---------------------------------  -----------------------------  -------------------  --------------------
CustomOp/linalg::matrixMulMatrix0  contrib.epu.quadric_custom_op  [-4902.8f, 4902.8f]                    18

2026-06-19 03:53 - INFO - epu - codegen - START====================build_cpu_runnable_fx_relay
2026-06-19 03:53 - INFO - epu - codegen - START=======================quantize_to_chimera_fx
2026-06-19 03:53 - INFO - epu - codegen - START=================================relay_to_tir
2026-06-19 03:53 - INFO - epu - codegen - START===========================relay_to_epu_relay
2026-06-19 03:53 - INFO - epu - codegen - START==============================adapt_and_order
2026-06-19 03:53 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-06-19 03:53 - INFO - epu - codegen - START=============================plan_lrm_virtual
2026-06-19 03:53 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-06-19 03:53 - INFO - epu - codegen - START===============================lrm_alloc_loop
2026-06-19 03:53 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-06-19 03:53 - INFO - epu - codegen - START================================lrm_splitting
2026-06-19 03:53 - INFO - epu - codegen - START==============================ext_split_relay
2026-06-19 03:53 - INFO - epu - codegen - START====================================build_tir
2026-06-19 03:53 - INFO - epu - chimera_job - Compilation of matmulnbits_single_layer_replaced_QC_N_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 successful



╒═════════════════════╤══════════════════════════════════════════════════════════════════════════════════╕
│ Module Name         │ matmulnbits_single_layer_replaced_QC_N_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ ONNX File           │ ./matmulnbits_single_layer_replaced.onnx                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Product Target      │ QC-N                                                                             │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Number of Cores     │ 1                                                                                │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ ISS Clock Frequency │ 1.700                                                                            │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ L2M Size            │ 16MB                                                                             │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ LRM Size            │ 4kB                                                                              │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ External Read BW    │ 128GBps                                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ External Write BW   │ 128GBps                                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ MACS per PE         │ 16                                                                               │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Max L2M             │ 0.000MB                                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Max LRM             │ 0.000kB                                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Max Temp Ext Bytes  │ 0.004MB                                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────────────────────────┤
│ Network GMACs       │                                                                                  │
╘═════════════════════╧══════════════════════════════════════════════════════════════════════════════════╛

╒════╤════════╤════════╤══════════════╤══════════════════════════╤═══════╕
│    │ Type   │ Name   │ shape        │ type                     │ mse   │
╞════╪════════╪════════╪══════════════╪══════════════════════════╪═══════╡
│  0 │ Input  │ A      │ [1, 1, 384]  │ tensor[FixedPoint32<27>] │ n/a   │
├────┼────────┼────────┼──────────────┼──────────────────────────┼───────┤
│  1 │ Output │ Y      │ [1, 1, 1152] │ tensor[FixedPoint32<18>] │ n/a   │
╘════╧════════╧════════╧══════════════╧══════════════════════════╧═══════╛



2026-06-19 03:53 - INFO - epu - iss_testing - Found tranges for input: <tvm.contrib.epu.interval.Interval object at 0x7651f47cb430>
2026-06-19 03:53 - INFO - epu - iss_testing - Started Executing Onnxruntime...
2026-06-19 03:53 - INFO - epu - chimera_job - running ort reference with overriden model: ./matmulnbits_single_layer_ort_fixed.onnx
2026-06-19 03:53 - INFO - epu - iss_testing - Done 0:00:00.075431
2026-06-19 03:53 - INFO - epu - iss_testing - Found tranges for input: <tvm.contrib.epu.interval.Interval object at 0x7651f4717010>
FILM 1/1: 100%|███████████████████████████████████████████████████████| 1/1 [00:00<00:00, 12.91it/s]



=== ORT vs ISS Comparison ===
ORT output range: [-2128.5337, 2209.0381]
ISS output range: [-2128.5337, 2209.0378]
RMSE:          0.0001
Correlation:   1.000000
Max abs error: 0.0002