add-inference-backend

Adding a New Inference Backend to AutoRound

Overview

This skill guides you through adding a new inference backend for running quantized models on a specific hardware platform. A backend defines how quantized weights are unpacked and computed at inference time. AutoRound automatically selects the best available backend based on hardware, quantization config, and priority.

Prerequisites

Before starting, determine:

1. Target hardware: CPU (Intel/AMD), CUDA GPU, Intel XPU, Habana HPU, etc.
2. Supported quantization configs: Which bit-widths, group sizes, and data types your backend handles
3. Kernel implementation: Triton, CUDA C++, PyTorch native, or external library (e.g., GPTQModel Marlin)
4. Packing format: How quantized weights are stored in memory

Step 1: Register Backend Info

Edit auto_round/inference/backend.py to register your backend's capabilities:

BackendInfos["auto_round:your_backend"] = BackendInfo(
    device=["cuda"],  # Supported devices
    sym=[True, False],  # Symmetric and/or asymmetric
    packing_format=["auto_round"],  # Compatible packing formats
    bits=[2, 4, 8],  # Supported bit-widths
    group_size=[32, 64, 128, -1],  # Supported group sizes (-1 = per-channel)
    compute_dtype=["float16", "bfloat16"],  # Compute precision
    data_type=["int"],  # Quantization data types
    act_bits=[16, 32],  # Activation bit-widths (16 = WxA16)
    priority=2,  # Higher = preferred (0-5 typical range)
    checkers=[your_feature_checker],  # Validation functions (optional)
    alias=["your_backend_short"],  # Alternative names (optional)
    requirements=["some_package>=1.0"],  # Required packages (optional)
    systems=["linux"],  # OS restriction (optional)
)

BackendInfo Fields Reference

Checker Functions

Use these pre-built checkers or create your own:

# Require in_features and out_features divisible by 32
from auto_round.inference.backend import feature_multiply_checker_32

# Require in_features divisible by group_size
from auto_round.inference.backend import in_feature_checker_group_size


# Custom checker
def your_feature_checker(in_feature, out_feature, config):
    """Check if layer dimensions are compatible with your backend."""
    return in_feature % 64 == 0 and out_feature % 64 == 0 and config["group_size"] in [64, 128]

Step 2: Implement QuantLinear Module

Create auto_round_extension/your_device/qlinear_your_backend.py:

import torch
import torch.nn as nn

QUANT_TYPE = "your_backend"


class QuantLinear(nn.Module):
    """Quantized linear layer for your backend.

    Stores packed quantized weights and performs dequantize-then-matmul
    (or fused quantized matmul) at inference time.
    """

    QUANT_TYPE = QUANT_TYPE

    def __init__(self, bits, group_size, in_features, out_features, bias=True, sym=True, **kwargs):
        super().__init__()
        self.bits = bits
        self.group_size = group_size
        self.in_features = in_features
        self.out_features = out_features
        self.sym = sym

        # Register packed weight buffers
        # Example: INT4 packed into INT32
        pack_factor = 32 // bits
        self.register_buffer(
            "qweight",
            torch.zeros(in_features // pack_factor, out_features, dtype=torch.int32),
        )
        self.register_buffer(
            "scales",
            torch.zeros(
                (in_features // group_size, out_features),
                dtype=torch.float16,
            ),
        )
        if not sym:
            self.register_buffer(
                "qzeros",
                torch.zeros(
                    (in_features // group_size, out_features // pack_factor),
                    dtype=torch.int32,
                ),
            )
        if bias:
            self.register_buffer("bias", torch.zeros(out_features, dtype=torch.float16))
        else:
            self.bias = None

    def forward(self, x):
        """Dequantize weights and compute linear transformation."""
        weight = self._dequantize()
        out = torch.matmul(x, weight.T)
        if self.bias is not None:
            out += self.bias
        return out

    def _dequantize(self):
        """Unpack and dequantize weights."""
        # Implement your dequantization kernel here
        # Can use Triton, CUDA, or PyTorch operations
        ...

    @classmethod
    def pack(cls, linear, scales, zeros, bias=None):
        """Pack a standard nn.Linear into this quantized format.

        Called during export to convert calibrated weights into packed format.
        """
        ...

Step 3: Wire Up QuantLinear Import Logic

Register your backend in the explicit import logic in auto_round/inference/backend.py. In this repository, backend loading is not a generic directory scan; dynamic_import_inference_linear() maps backend keys to specific QuantLinear implementations.

Add a new backend key in BackendInfos[...] if needed, and make sure dynamic_import_inference_linear() returns your QuantLinear class for that backend:

if backend == "auto_round:your_backend":
    from auto_round_extension.your_device.qlinear_your_backend import QuantLinear

    return QuantLinear

If your backend fits an existing branch pattern, you can also reuse that logic, but contributors should update the explicit import mapping rather than rely on implicit auto-discovery.

Step 4: Add Extension __init__.py

Create auto_round_extension/your_device/__init__.py if the directory is new:

# Auto-Round extension for YourDevice backend

Step 5: Test

Unit test for the QuantLinear

def test_your_backend_qlinear():
    from auto_round_extension.your_device.qlinear_your_backend import QuantLinear

    ql = QuantLinear(bits=4, group_size=128, in_features=256, out_features=512)
    x = torch.randn(1, 256, dtype=torch.float16, device="cuda")
    out = ql(x)
    assert out.shape == (1, 512)

End-to-end test

def test_your_backend_e2e(tiny_opt_model_path, dataloader):
    ar = AutoRound(
        tiny_opt_model_path,
        bits=4,
        group_size=128,
        dataset=dataloader,
        iters=2,
        nsamples=2,
    )
    compressed_model, _ = ar.quantize()
    ar.save_quantized(output_dir="./tmp_backend_test", format="auto_round")

    # Load and verify inference with your backend
    from transformers import AutoModelForCausalLM, AutoTokenizer

    model = AutoModelForCausalLM.from_pretrained("./tmp_backend_test")
    tokenizer = AutoTokenizer.from_pretrained("./tmp_backend_test")
    inputs = tokenizer("Hello", return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=10)
    assert outputs.shape[1] > inputs["input_ids"].shape[1]

Reference: Existing Backend Implementations

Key Registration Points