Add Inference Backend

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)
)

# 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]

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.
        """
        ...