Blame for onnxruntime/quantization/matmul_nbits_quantizer.py

# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------

from __future__ import annotations

import argparse
import copy
import logging
import os

import ml_dtypes
import numpy as np
import numpy.typing as npt
import onnx
import onnx_ir as ir
from onnx.onnx_pb import GraphProto, ModelProto, NodeProto, TensorProto

from onnxruntime.capi._pybind_state import (
    quantize_matmul_2bits,
    quantize_matmul_4bits,
    quantize_matmul_8bits,
    quantize_qdq_matmul_4bits,
)

from .calibrate import CalibrationDataReader
from .neural_compressor import gptq_quantize, rtn_quantize
from .onnx_model import ONNXModel
from .quant_utils import QuantFormat, attribute_to_kwarg

logging.basicConfig(format="%(asctime)s %(name)s [%(levelname)s] - %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)


class WeightOnlyQuantConfig:
    def __init__(
        self,
        algorithm: str,
        quant_format: QuantFormat,
        op_types_to_quantize: tuple[str, ...] | None = None,
        quant_axes: tuple[tuple[str, int], ...] | None = None,
        customized_weight_config: dict | None = None,
    ):
        """This is the Base class for Weight Only blockwise quantization Configuration.

        Args:
            algorithm:
                weight only quantize algorithm name.
            quant_format: QuantFormat{QOperator, QDQ}.
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
            op_types_to_quantize (optional):
                set of operator types to quantize. Default {MatMul}
            quant_axes (dict[str, int], optional):
                op:axis, which axis to quantize for an op. Default {MatMul: 0, Gather: 1}
            customized_weight_config:
                customized weight config for nodes if needed. It is dictionary with node name as key,
                and the value is a dict of customized config.
        """
        self.algorithm = algorithm
        self.quant_format = quant_format
        self.op_types_to_quantize = set(op_types_to_quantize) if op_types_to_quantize else {"MatMul"}
        self.quant_axes = dict(quant_axes) if quant_axes else {"MatMul": 0, "Gather": 1}
        self.customized_weight_config = customized_weight_config


class RTNWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        ratios=None,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
        customized_weight_config: dict | None = None,
    ):
        """
        This is a class for round-to-nearest (RTN) algorithm Weight Only Quant Configuration.
        RTN is the most straightforward way to quantize weight using scale maps.

        Args:
            ratios:
                percentile of clip. Defaults to {}.
            quant_format (QuantFormat{QOperator, QDQ}, optional):
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
                Defaults to QuantFormat.QOperator.
            op_types_to_quantize (optional):
                set of operator types to quantize.
            customized_weight_config:
                customized weight config for nodes if needed. It is dictionary with node name as key,
                and the value is a dict of customized config.
        """
        assert quant_format == QuantFormat.QOperator, "RTN only supports QOperator format"

        if ratios is None:
            ratios = {}
        super().__init__(
            algorithm="RTN",
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
            customized_weight_config=customized_weight_config,
        )
        self.ratios = ratios


class KQuantWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        ratios=None,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
        customized_weight_config: dict | None = None,
    ):
        """
        This is a class for k-quant algorithm Weight Only Quant Configuration.

        Args:
            ratios:
                percentile of clip. Defaults to {}.
            quant_format (QuantFormat{QOperator, QDQ}, optional):
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
                Defaults to QuantFormat.QOperator.
            op_types_to_quantize (optional):
                set of operator types to quantize.
        """
        assert quant_format == QuantFormat.QOperator, "k-quant only supports QOperator format"

        if ratios is None:
            ratios = {}
        super().__init__(
            algorithm="k_quant",
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
            customized_weight_config=customized_weight_config,
        )
        self.ratios = ratios


class GPTQWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        calibration_data_reader: CalibrationDataReader | None = None,
        percdamp=0.01,
        block_size=128,
        actorder=False,
        mse=False,
        perchannel=True,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
    ):
        """
        This is a class for GPTQ algorithm Weight Only Quant Configuration.
        GPTQ algorithm provides more accurate quantization but requires more computational resources.

        Args:
            calibration_data_reader:
                a calibration data reader. It enumerates calibration data and generates inputs for the original model.
            percdamp:
                percent of the average Hessian diagonal to use for dampening.
            block_size (int, optional):
                channel number in one block to execute a GPTQ quantization iteration.
            actorder (bool, optional):
                whether rearrange Hessian matrix considering the diag's value.
            mse (bool, optional):
                whether get scale and zero point with mse error.
            perchannel (bool, optional):
                whether quantize weight per-channel.
            quant_format (QuantFormat{QOperator, QDQ}, optional):
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
                Defaults to QuantFormat.QOperator.
            op_types_to_quantize (optional):
                set of operator types to quantize.
        """
        assert quant_format == QuantFormat.QOperator, "GPTQ only supports QOperator format"

        super().__init__(
            algorithm="GPTQ",
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
        )
        self.calibration_data_reader = calibration_data_reader
        self.percdamp = percdamp
        self.block_size = block_size
        self.actorder = actorder
        self.mse = mse
        self.perchannel = perchannel


class HQQWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        block_size=128,
        bits=4,
        axis=1,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
        quant_axes: tuple[tuple[str, int], ...] | None = None,
    ):
        """
        This is a class for HQQ algorithm Weight Only Quant Configuration.
        HQQ algorithm quant weight without needing calibrate data.

        Args:
            block_size (int, optional):
                channel number in one block to execute a HQQ quantization iteration.
            bits (int, optional):
                how many bits to represent weight.
            axis (int, optional):
                0 or 1. which axis to quantize. https://arxiv.org/pdf/2309.15531.pdf
            quant_format (QuantFormat{QOperator, QDQ}, optional):
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
                Defaults to QuantFormat.QOperator.
            op_types_to_quantize (optional):
                set of operator types to quantize.
            quant_axes (dict[str, int], optional):
                op:axis, which axis to quantize for an op. Default {MatMul: 0, Gather: 1}
        """
        assert quant_format == QuantFormat.QOperator, "HQQ only supports QOperator format"

        super().__init__(
            algorithm="HQQ",
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
            quant_axes=quant_axes,
        )
        self.block_size = block_size
        self.bits = bits
        self.axis = axis


class DefaultWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        block_size: int = 128,
        is_symmetric: bool = False,
        accuracy_level: int | None = None,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
        quant_axes: tuple[tuple[str, int], ...] | None = None,
        bits: int = 4,
        channel_wised_quantize: bool = False,
    ):
        """
        This is a class for weight only affine quantization configuration.

        Args:
            block_size (int, optional):
                channel number in one block to execute an affine quantization iteration.
            is_symmetric (bool, optional):
                whether quantize weight symmetrically.
            accuracy_level (int, optional):
                Accuracy level of the 4-bit quantized MatMul computation.
                Refer to the MatMulNBits contrib op's 'accuracy_level' attribute for details.
                (https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits)
            quant_format (QuantFormat{QOperator, QDQ}, optional):
                QOperator format quantizes the model with quantized operators directly.
                QDQ format quantize the model by inserting QuantizeLinear/DeQuantizeLinear on the tensor.
                Defaults to QuantFormat.QOperator.
            op_types_to_quantize (optional):
                set of operator types to quantize.
            quant_axes (dict[str, int], optional):
                op:axis, which axis to quantize for an op. Default {MatMul: 0, Gather: 1}
            bits (int, optional):
                number of bits per element after quantization. Default 4.
        """
        super().__init__(
            algorithm="DEFAULT",
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
            quant_axes=quant_axes,
        )
        self.block_size = block_size
        self.is_symmetric = is_symmetric
        self.bits = bits
        self.accuracy_level = accuracy_level
        self.channel_wised_quantize = channel_wised_quantize
        if channel_wised_quantize and quant_format == QuantFormat.QOperator:
            raise NotImplementedError("QuantFormat.QOperator is not supported channel_wised_quantize yet")


class NVAWQWeightOnlyQuantConfig(WeightOnlyQuantConfig):
    def __init__(
        self,
        tokenizer_dir,
        dataset_name="cnn",
        cache_dir="./cache",
        calibration_method="awq_lite",
    ):
        """
        Configuration for the nvidia_awq quantization method.

        Args:
            tokenizer_dir (str): pathof the tokenizer dir.
            dataset_name (str): Name of the dataset.
            cache_dir (str): Directory for caching.
            calibration_method (str): calib method for nvidia_awq.
        """
        # Import torch and DataLoader
        try:
            import torch  # noqa: PLC0415
            from torch.utils.data import DataLoader  # noqa: PLC0415

            self.torch = torch
            self.DataLoader = DataLoader
        except ImportError:
            print(
                "Error: The 'torch' library is required but not installed. Please install it using 'pip install torch'."
            )
            raise ImportError("torch is not installed. Exiting.") from None

        # Import datasets
        try:
            from datasets import load_dataset  # noqa: PLC0415

            self.load_dataset = load_dataset
        except ImportError:
            print(
                "Error: The 'datasets' library is required but not installed. Please install it using 'pip install datasets'."
            )
            raise ImportError("datasets is not installed. Exiting.") from None

        # Import transformers
        try:
            from transformers import AutoConfig, AutoTokenizer  # noqa: PLC0415

            self.AutoConfig = AutoConfig
            self.AutoTokenizer = AutoTokenizer
        except ImportError:
            print(
                "Error: The 'transformers' library is required but not installed. Please install it using 'pip install transformers'."
            )
            raise ImportError("transformers is not installed. Exiting.") from None

        super().__init__(
            algorithm="nvidia_awq",
            quant_format=QuantFormat.QDQ,
            op_types_to_quantize=None,  # Assuming op_types_to_quantize is handled elsewhere
            quant_axes=None,  # Assuming quant_axes is handled elsewhere
        )

        # Determine the device
        device = self.torch.device("cuda" if self.torch.cuda.is_available() else "cpu")

        calib_inputs = self.get_calib_inputs(
            dataset_name=dataset_name,
            model_name=tokenizer_dir,
            cache_dir=cache_dir,
            calib_size=32,
            batch_size=1,
            block_size=512,
            device=device,
            use_fp16=True,
            use_buffer_share=False,
            add_past_kv_inputs=True,
            max_calib_rows_to_load=128,
            add_position_ids=True,
        )

        self.calibration_data_reader = calib_inputs
        self.calibration_method = calibration_method

    def make_model_input(
        self,
        config,
        input_ids_arg,
        attention_mask_arg,
        add_past_kv_inputs,
        device,
        use_fp16,
        use_buffer_share,
        add_position_ids,
    ):
        # Access torch from the instance variable
        torch = self.torch

        input_ids = input_ids_arg
        attention_mask = attention_mask_arg

        if isinstance(input_ids_arg, list):
            input_ids = torch.tensor(input_ids_arg, device=device, dtype=torch.int64)
            attention_mask = torch.tensor(attention_mask_arg, device=device, dtype=torch.int64)

        inputs = {
            "input_ids": input_ids.contiguous(),
            "attention_mask": attention_mask.contiguous(),
        }

        if add_position_ids:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            inputs["position_ids"] = position_ids.contiguous()

        if add_past_kv_inputs:
            torch_dtype = torch.float16 if use_fp16 else torch.float32
            batch_size, sequence_length = input_ids.shape
            max_sequence_length = config.max_position_embeddings
            num_heads, head_size = (
                config.num_key_value_heads,
                config.hidden_size // config.num_attention_heads,
            )
            for i in range(config.num_hidden_layers):
                past_key = torch.zeros(
                    batch_size,
                    num_heads,
                    max_sequence_length if use_buffer_share else 0,
                    head_size,
                    device=device,
                    dtype=torch_dtype,
                )
                past_value = torch.zeros(
                    batch_size,
                    num_heads,
                    max_sequence_length if use_buffer_share else 0,
                    head_size,
                    device=device,
                    dtype=torch_dtype,
                )
                inputs.update(
                    {
                        f"past_key_values.{i}.key": past_key.contiguous(),
                        f"past_key_values.{i}.value": past_value.contiguous(),
                    }
                )

        return inputs

    def get_calib_inputs(
        self,
        dataset_name,
        model_name,
        cache_dir,
        calib_size,
        batch_size,
        block_size,
        device,
        use_fp16,
        use_buffer_share,
        add_past_kv_inputs,
        max_calib_rows_to_load,
        add_position_ids,
    ):
        # Access transformers and datasets from the instance variables
        auto_config = self.AutoConfig
        auto_tokenizer = self.AutoTokenizer
        load_dataset = self.load_dataset

        config = auto_config.from_pretrained(
            model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=True
        )
        tokenizer = auto_tokenizer.from_pretrained(
            model_name, use_auth_token=True, cache_dir=cache_dir, trust_remote_code=True
        )
        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
        tokenizer.pad_token = tokenizer.eos_token

        assert calib_size <= max_calib_rows_to_load, "calib size should be no more than max_calib_rows_to_load"

        if "cnn" in dataset_name:
            dataset2 = load_dataset("cnn_dailymail", name="3.0.0", split="train").select(range(max_calib_rows_to_load))
            column = "article"
        elif "pile" in dataset_name:
            dataset2 = load_dataset("mit-han-lab/pile-val-backup", split="validation")
            column = "text"
        else:
            raise ValueError(f'dataset "{dataset_name}" not supported')

        dataset2 = dataset2[column][:calib_size]
        batch_encoded = tokenizer.batch_encode_plus(
            dataset2, return_tensors="pt", padding=True, truncation=True, max_length=block_size
        )
        batch_encoded = batch_encoded.to(device)
        batch_encoded_input_ids = batch_encoded["input_ids"]
        batch_encoded_attention_mask = batch_encoded["attention_mask"]

        # Access DataLoader from the instance variable
        data_loader = self.DataLoader

        calib_dataloader_input_ids = data_loader(batch_encoded_input_ids, batch_size=batch_size, shuffle=False)
        calib_dataloader_attention_mask = data_loader(
            batch_encoded_attention_mask, batch_size=batch_size, shuffle=False
        )

        assert len(calib_dataloader_input_ids.dataset) == len(calib_dataloader_attention_mask.dataset)
        assert len(calib_dataloader_input_ids) == len(calib_dataloader_attention_mask)

        number_of_batched_samples = calib_size // batch_size

        batched_input_ids = []
        for idx, data in enumerate(calib_dataloader_input_ids):
            batched_input_ids.append(data)
            if idx == (number_of_batched_samples - 1):
                break

        batched_attention_mask = []
        for idx, data in enumerate(calib_dataloader_attention_mask):
            batched_attention_mask.append(data)
            if idx == (number_of_batched_samples - 1):
                break

        print(
            f"\n--Quantize-Script-- number_of_batched_samples={number_of_batched_samples}, "
            f"batch-input-ids-list-len={len(batched_input_ids)}, batched_attention_mask={len(batched_attention_mask)}\n"
        )

        batched_inputs_list = []
        for i in range(number_of_batched_samples):
            input_ids = batched_input_ids[i]
            attention_mask = batched_attention_mask[i]

            inputs = self.make_model_input(
                config,
                input_ids,
                attention_mask,
                add_past_kv_inputs,
                device,
                use_fp16,
                use_buffer_share,
                add_position_ids,
            )
            inputs = {input_name: torch_tensor.cpu().numpy() for input_name, torch_tensor in inputs.items()}
            batched_inputs_list.append(inputs)

        print(f"\n--Quantize-Script-- number of batched inputs = {len(batched_inputs_list)}\n")
        return batched_inputs_list


def is_divisible(val1, val2):
    return int(val2 * np.ceil(val1 / val2)) == val1


class HQQWeightOnlyQuantizer:
    def __init__(
        self,
        config: HQQWeightOnlyQuantConfig,
    ):
        self.config = config

    # Proximal solver || weight - dequantize(quantize(weight))||_p^p
    @staticmethod
    def optimize_weights(
        tensor,
        scale,
        zero,
        min_max: list[int],
        axis: int = 0,
        opt_params: dict | None = None,
        verbose=False,
    ):
        import torch  # noqa: PLC0415

        opt_params = {"lp_norm": 0.7, "beta": 1e1, "kappa": 1.01, "iters": 20} if opt_params is None else opt_params
        lp_norm, beta, kappa, iters = (
            opt_params["lp_norm"],
            opt_params["beta"],
            opt_params["kappa"],
            opt_params["iters"],
        )

        dtype = torch.float16 if tensor.is_cuda else torch.float32
        w_f = tensor.to(dtype)
        scale = scale.to(dtype)
        zero = zero.to(dtype)

        def shrink_op(x, beta, p=lp_norm):
            if p == 1:
                return torch.sign(x) * torch.nn.functional.relu(torch.abs(x) - 1.0 / beta)
            else:
                return torch.sign(x) * torch.nn.functional.relu(
                    torch.abs(x) - (1.0 / beta) * torch.pow(torch.abs(x) + 1e-8, p - 1)
                )

        best_error = 1e4
        for i in range(iters):
            w_q = torch.round(w_f * scale + zero).clamp(min_max[0], min_max[1])
            w_r = (w_q - zero) / scale
            w_e = shrink_op(w_f - w_r, beta)
            zero = torch.mean(w_q - (w_f - w_e) * scale, axis=axis, keepdim=True)
            beta *= kappa

            current_error = float(torch.abs(w_f - w_r).mean())
            if verbose:
                print(i, np.round(current_error, 6))
            if current_error < best_error:
                best_error = current_error
            else:
                break

        del w_f, w_q, w_r, w_e

        return scale, zero

    @staticmethod
    def pack_on_row_fast_248bit(pack_tensor, ori_int_tensor, bits):
        if pack_tensor.shape[0] == ori_int_tensor.shape[0]:
            ori_int_tensor = ori_int_tensor.T
            pack_tensor = pack_tensor.T
        if bits in [2, 4, 8]:
            compress_ratio = pack_tensor.element_size() * 8 // bits
            for j in range(compress_ratio):
                pack_tensor[0:] |= ori_int_tensor[j::compress_ratio] << (bits * (j))
        else:
            raise NotImplementedError("Only 2,4,8 bits are supported.")

    # from Official implementation of Half-Quadratic Quantization (HQQ)
    def quantize_internal(
        self, tensor, bits=4, channel_wise=True, group_size=64, optimize=True, round_zero=True, axis=1
    ):
        import torch  # noqa: PLC0415

        weight = tensor.float()
        ori_shape = weight.shape

        pad_len = (group_size - ori_shape[axis] % group_size) % group_size
        if axis == 1:
            weight = torch.nn.functional.pad(weight, (0, pad_len), "constant", 0)
        else:
            weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_len), "constant", 0)
        shape = weight.shape

        # Reshape for grouping
        if (group_size is not None) and channel_wise:
            weight = weight.reshape([-1, group_size]) if (axis == 1) else weight.reshape([group_size, -1])

        # Get min/max values
        if channel_wise is False:
            _min, _max = weight.min(), weight.max()
            optimize = False
        else:
            _min = weight.min(axis=axis, keepdim=True)[0]
            _max = weight.max(axis=axis, keepdim=True)[0]

        max_v = 2**bits - 1
        min_v = 0
        min_max = [min_v, max_v]

        # Note: here we work with the inverse of the scale to avoid division and quantize instead via weight*scale + zero, the scale is inverted later on.
        # clamp to avoid half-precision problems
        scale = (max_v / (_max - _min)).clamp(max=2e4)
        #!!!!!!!!!!!!!!!
        min_max_axis = _max - _min
        if (min_max_axis == 0).sum().item() > 0:
            min_max_axis[min_max_axis == 0] = max_v
            scale = (max_v / min_max_axis).clamp(max=2e4)
        zero = -_min * scale

        if round_zero:
            zero = torch.round(zero)

        # Fine-tune weights
        if optimize:
            scale, zero = self.optimize_weights(tensor=weight, scale=scale, zero=zero, min_max=min_max, axis=axis)

        # Quantize
        # Necessary for fake quantization backprop
        w_q = torch.round(weight * scale + zero).clamp(min_max[0], min_max[1])
        w_q = w_q.reshape(shape).int()

        scale = 1.0 / scale
        if axis == 1:
            scale = scale.reshape(shape[0], -1)
            zero = zero.reshape(shape[0], -1)
        else:
            scale = scale.reshape(-1, shape[-1])
            zero = zero.reshape(-1, shape[-1])
        # cleanup
        del weight, _min, _max

        return w_q, scale.to(tensor.dtype), zero.to(tensor.dtype)

    def quantize(self, node: NodeProto, graph_stack: list[GraphProto]) -> list[NodeProto]:
        """
        Target node:        QOperator node:            QDQ nodes:
        MatMul              MatMulNBits                DeQuantizeLinear -> MatMul
        Gather              GatherBlockQuantized       Gather, Gather, Gather (optional) -> DequantizeLinear
        If the node is target node with fp32 or fp16 const weight, quantize the weight to int4 and
        return the new nodes.
        If QOperator format, return the corresponding QOperator nodes.
        If QDQ format, return the corresdponging QDQ nodes.
        Gather (quantized data) + Gather (scales) + Gather (optional, zero points) -> DequantizeLinear is
        not supported yet because Gather does not support int4 data.
        """
        # With HQQ, zero points are in float. Current GatherBlockQuantized does not support float zero points.
        if node.op_type == "Gather":
            raise NotImplementedError("Gather quantization is not supported yet in HQQ")

        import torch  # noqa: PLC0415

        logger.info(f"start to quantize {node.name} ...")
        input_b = node.input[1]
        b_pb, bs_graph = get_initializer(input_b, graph_stack)
        if b_pb is None:
            logger.info("MatMul doesn't have const weight. Skip to quantize")
            return [node]  # only care about constant weight

        b_array = onnx.numpy_helper.to_array(b_pb)
        b_original_shape = b_array.shape
        if len(b_original_shape) != 2:
            if len(b_original_shape) < 2:
                logger.info("MatMul weight has fewer than 2 dimensions. Skip to quantize.")
                return [node]
            leading = b_original_shape[:-2]
            if any(d != 1 for d in leading):
                logger.info(
                    "MatMul weight has non-unit batch dims %s; N-D batched quantization not supported. "
                    "Skip to quantize.",
                    list(leading),
                )
                return [node]
            # Squeeze all unit leading dims to obtain a 2-D weight [K, N]
            b_array = b_array.reshape(b_original_shape[-2], b_original_shape[-1])
        else:
            b_original_shape = None  # already 2-D, no reshape needed
        b_array_torch = torch.from_numpy(b_array)
        if torch.cuda.is_available():
            b_array_torch = b_array_torch.cuda()

        bits = self.config.bits
        quant_weight_torch, scales_torch, zero_points_torch = self.quantize_internal(
            b_array_torch.T, bits=bits, group_size=self.config.block_size
        )
        quant_weight_torch = quant_weight_torch.contiguous()
        scales_torch = scales_torch.contiguous()
        zero_points_torch = zero_points_torch.contiguous()

        packed_size = 8 // bits  # number of elements packed into one byte

        packed_torch = torch.zeros(
            (quant_weight_torch.shape[0], quant_weight_torch.shape[1] // packed_size),
            dtype=torch.uint8,
            device=quant_weight_torch.device,
        )
        self.pack_on_row_fast_248bit(packed_torch, quant_weight_torch, bits)
        scales = scales_torch.cpu().numpy()
        zero_points = zero_points_torch.cpu().numpy()
        # reshape to the predefined shape in MatmulNbits
        scales = scales.reshape(-1)
        zero_points = zero_points.reshape(-1)
        rows, cols = b_array_torch.shape
        block_size = self.config.block_size
        blob_size = block_size // packed_size
        k_blocks = (rows + block_size - 1) // block_size
        packed_torch = packed_torch.reshape(cols, k_blocks, blob_size)

        b_quant = onnx.numpy_helper.from_array(packed_torch.cpu().numpy())
        b_quant.name = b_pb.name + "_Q" + str(bits)
        for input in bs_graph.input:
            if input.name == input_b:
                bs_graph.input.remove(input)
                break

        scales_tensor = onnx.numpy_helper.from_array(scales)
        scales_tensor.name = b_pb.name + "_scales"
        bs_graph.initializer.extend([b_quant, scales_tensor])

        input_names = [node.input[0], b_quant.name, scales_tensor.name]
        zp_tensor = onnx.numpy_helper.from_array(zero_points)
        zp_tensor.name = b_pb.name + "_zero_points"
        bs_graph.initializer.extend([zp_tensor])
        input_names.append(zp_tensor.name)

        kwargs = {}
        rows, cols = b_array.shape
        kwargs["K"] = rows
        kwargs["N"] = cols
        kwargs["bits"] = bits
        kwargs["block_size"] = self.config.block_size

        # When rank(A) >= rank_b_orig, the reshape would be a no-op since MatMulNBits
        # already produces the correct output shape.
        a_static_rank = _get_static_rank(node.input[0], graph_stack)
        rank_b_orig = len(b_original_shape) if b_original_shape is not None else 0
        needs_reshape = b_original_shape is not None and (a_static_rank is None or a_static_rank < rank_b_orig)
        matmul_q_output = node.output[0] if not needs_reshape else node.output[0] + "_pre_reshape"
        matmul_q_node = onnx.helper.make_node(
            "MatMulNBits",
            inputs=input_names,
            outputs=[matmul_q_output],
            name=node.name + "_Q" + str(bits) if node.name else "",
            domain="com.microsoft",
            **kwargs,
        )

        output_nodes = [matmul_q_node]
        if needs_reshape:
            # Restore ONNX MatMul broadcast shape on the output.
            # MatMul(A, B_orig) output shape (with B_orig leading dims all 1) is:
            #   [1] * max(rank(B_orig) - rank(A), 0) + A.shape[:-1] + [N]
            # MatMulNBits with squeezed B produces [...A.shape[:-1], N] (rank=rank(A)),
            # so we only need to prepend leading 1s when rank(B_orig) > rank(A).
            output_nodes.extend(
                _build_nbits_output_reshape(
                    a_input_name=node.input[0],
                    b_original_shape=b_original_shape,
                    target_graph=bs_graph,
                    name_prefix=(node.name + "_Q" + str(bits)) if node.name else (b_pb.name + "_Q" + str(bits)),
                    pre_reshape_output=matmul_q_output,
                    final_output=node.output[0],
                )
            )

        logger.info(f"complete quantization of {node.name} ...")

        return output_nodes


def get_initializer(name, graph_path: list[GraphProto]) -> tuple[TensorProto, GraphProto]:
    for gid in range(len(graph_path) - 1, -1, -1):
        graph = graph_path[gid]
        for tensor in graph.initializer:
            if tensor.name == name:
                return tensor, graph
    return None, None


def _get_static_rank(tensor_name: str, graph_path: list[GraphProto]) -> int | None:
    """Return the static rank of a tensor if its shape is known, else None.

    Searches graph inputs, value_info, and outputs in the graph stack (inner-most
    graph first).  A known shape requires ``HasField('shape')`` to be true on the
    tensor_type; the rank is then ``len(shape.dim)``.  Individual dim sizes may
    still be symbolic — only the rank (dimension count) matters here.
    """
    for gid in range(len(graph_path) - 1, -1, -1):
        graph = graph_path[gid]
        for vi in list(graph.input) + list(graph.value_info) + list(graph.output):
            if vi.name == tensor_name:
                tt = vi.type.tensor_type
                if tt.HasField("shape"):
                    return len(tt.shape.dim)
                return None
    return None


def _build_nbits_output_reshape(
    a_input_name: str,
    b_original_shape: tuple,
    target_graph: GraphProto,
    name_prefix: str,
    pre_reshape_output: str,
    final_output: str,
) -> list[NodeProto]:
    """Build the reshape chain that restores the ONNX MatMul-broadcast output shape.

    MatMulNBits produces shape ``[...A_batch_dims, M, N]`` (rank = rank(A)). To match
    the original ``MatMul(A, B_orig)`` output, where B_orig has all-unit leading
    dims, we need:

        a_rank_eff = max(rank(A), 2)   # ONNX promotes 1-D A to rank-2
        out_shape = [1] * max(rank(B_orig) - a_rank_eff, 0) + A.shape[:-1] + [N]

    This is built dynamically via Shape/Gather/Max/Sub/Max/ConstantOfShape/Slice/Concat
    so it works regardless of A's static rank (handles rank(A) == 1, rank(A) == 2
    — the common transformer case — as well as rank(A) >= rank(B_orig) where no
    leading-1 prepending is needed). All ops used are valid from opset 11 onward.

    Args:
        a_input_name: name of the activation input edge (A) feeding MatMulNBits.
        b_original_shape: the original (pre-squeeze) shape of B, e.g. ``(1, K, N)``.
        target_graph: graph proto to append helper initializers into.
        name_prefix: unique prefix for generated node/initializer names.
        pre_reshape_output: name of the MatMulNBits output edge (the input of the
            generated Reshape).
        final_output: name of the final edge produced by the generated Reshape
            (must match the original MatMul output edge).

    Returns:
        List of nodes to append to the consumer's ``output_nodes`` after the
        MatMulNBits node. Initializers are appended to ``target_graph`` in place.
    """
    rank_b_orig = len(b_original_shape)
    n_dim = b_original_shape[-1]

    # Incorporate the unique output tensor name so initializer names are unique
    # even when multiple MatMul nodes share the same node.name (Fix 2).
    p = name_prefix + "_" + final_output
    init_zero = p + "_zero"
    init_one = p + "_one"
    init_two = p + "_two"
    init_one_vec = p + "_one_vec"
    init_rank_b = p + "_rank_b"
    init_n_vec = p + "_n_vec"
    init_zero_vec = p + "_zero_vec"
    init_one_value_template = p + "_one_value"

    target_graph.initializer.extend(
        [
            onnx.numpy_helper.from_array(np.array(0, dtype=np.int64), name=init_zero),
            onnx.numpy_helper.from_array(np.array(1, dtype=np.int64), name=init_one),
            onnx.numpy_helper.from_array(np.array(2, dtype=np.int64), name=init_two),
            onnx.numpy_helper.from_array(np.array([1], dtype=np.int64), name=init_one_vec),
            onnx.numpy_helper.from_array(np.array(rank_b_orig, dtype=np.int64), name=init_rank_b),
            onnx.numpy_helper.from_array(np.array([n_dim], dtype=np.int64), name=init_n_vec),
            onnx.numpy_helper.from_array(np.array([0], dtype=np.int64), name=init_zero_vec),
        ]
    )

    a_shape = p + "_a_shape"
    a_shape_of_shape = p + "_a_shape_of_shape"
    a_rank = p + "_a_rank"
    a_rank_eff = p + "_a_rank_eff"
    extra_raw = p + "_extra_raw"
    extra_count = p + "_extra_count"
    extra_count_vec = p + "_extra_count_vec"
    extra_ones = p + "_extra_ones"
    a_rank_minus_one = p + "_a_rank_m1"
    a_rank_minus_one_vec = p + "_a_rank_m1_vec"
    a_prefix_shape = p + "_a_prefix_shape"
    target_shape = p + "_target_shape"

    nodes = [
        onnx.helper.make_node("Shape", [a_input_name], [a_shape], name=p + "_shape_a"),
        # Use Shape(shape) + Gather instead of Size to stay within opset 11
        # (Size requires opset >= 13 when applied to a shape tensor).
        # Shape applied to the 1-D shape vector yields [rank_a] as a 1-element
        # tensor; Gather with scalar index 0 extracts it as a scalar int64.
        onnx.helper.make_node("Shape", [a_shape], [a_shape_of_shape], name=p + "_shape_of_a_shape"),
        onnx.helper.make_node("Gather", [a_shape_of_shape, init_zero], [a_rank], name=p + "_gather_rank"),
        # ONNX MatMul promotes a 1-D activation to rank-2 before computing the
        # output shape, so use Max(a_rank, 2) as the effective rank when
        # computing how many leading 1s to prepend.
        onnx.helper.make_node("Max", [a_rank, init_two], [a_rank_eff], name=p + "_max_rank_eff"),
        onnx.helper.make_node("Sub", [init_rank_b, a_rank_eff], [extra_raw], name=p + "_sub"),
        onnx.helper.make_node("Max", [extra_raw, init_zero], [extra_count], name=p + "_max"),
        onnx.helper.make_node("Reshape", [extra_count, init_one_vec], [extra_count_vec], name=p + "_reshape_extra"),
        onnx.helper.make_node(
            "ConstantOfShape",
            [extra_count_vec],
            [extra_ones],
            name=p + "_const_ones",
            value=onnx.helper.make_tensor(
                name=init_one_value_template,
                data_type=TensorProto.INT64,
                dims=[1],
                vals=[1],
            ),
        ),
        onnx.helper.make_node("Sub", [a_rank, init_one], [a_rank_minus_one], name=p + "_sub_one"),
        onnx.helper.make_node(
            "Reshape", [a_rank_minus_one, init_one_vec], [a_rank_minus_one_vec], name=p + "_reshape_rank_m1"
        ),
        onnx.helper.make_node(
            "Slice",
            [a_shape, init_zero_vec, a_rank_minus_one_vec, init_zero_vec],
            [a_prefix_shape],
            name=p + "_slice_a_prefix",
        ),
        onnx.helper.make_node(
            "Concat",
            [extra_ones, a_prefix_shape, init_n_vec],
            [target_shape],
            name=p + "_concat_target",
            axis=0,
        ),
        onnx.helper.make_node(
            "Reshape",
            [pre_reshape_output, target_shape],
            [final_output],
            name=p + "_reshape_out",
        ),
    ]
    return nodes


# transpose int4 matrix (packed as uint8)
def transpose_packed_int4_matrix(packed, rows, cols):
    # unpack to int4 matrix
    total = rows * cols
    high = (packed >> 4) & 0x0F
    low = packed & 0x0F
    int4_vals = np.empty(total, dtype=np.uint8)
    int4_vals[0::2] = low
    int4_vals[1::2] = high
    int4_matrix = int4_vals.reshape((rows, cols))

    # transpose int4 matrix
    int4_matrix_transposed = int4_matrix.T

    # pack to uint8
    flat = int4_matrix_transposed.reshape(-1)
    packed = ((flat[1::2] << 4) & 0xF0) | (flat[0::2] & 0x0F)
    return packed.astype(np.uint8)


class DefaultWeightOnlyQuantizer:
    def __init__(self, config: DefaultWeightOnlyQuantConfig):
        self.config = config

    def qbits_block_quant(self, fp32weight: npt.ArrayLike) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """4b/8b quantize fp32 weight to int4 using C++ kernels."""

        qbits = self.config.bits
        kpack = 8 // qbits
        if len(fp32weight.shape) != 2:
            raise ValueError("Current int4 block quantization only supports 2D tensors!")
        rows, cols = fp32weight.shape

        block_size = self.config.block_size
        k_blocks = (rows + block_size - 1) // block_size

        if self.config.quant_format == QuantFormat.QOperator:
            blob_size = (block_size + kpack - 1) // kpack
            padded_rows = k_blocks * block_size
            pad_len = padded_rows - rows
            if pad_len > 0:
                fp32weight = np.pad(fp32weight, ((0, pad_len), (0, 0)), "constant")

            # block wise quantization, each block comes from a single column
            packed = np.zeros((cols, k_blocks, blob_size), dtype="uint8")
            zero_point = np.zeros((cols, ((k_blocks + kpack - 1) // kpack)), dtype="uint8")
            scales = np.zeros((cols, k_blocks), dtype=fp32weight.dtype)
            if qbits == 2:
                quantize_matmul_2bits(
                    packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
                )
            elif qbits == 8:
                quantize_matmul_8bits(
                    packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
                )
            else:
                quantize_matmul_4bits(
                    packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
                )
        else:
            # block size equal to rows (K) if channel wised quantize enabled
            block_size = rows if self.config.channel_wised_quantize else self.config.block_size
            k_blocks = (rows + block_size - 1) // block_size

            assert qbits == 4, "QDQ format only support 4 bits quantization"
            packed = np.zeros((rows * cols + 1) // 2, dtype="uint8")
            zero_point = np.zeros((cols * k_blocks + 1) // 2, dtype="uint8")
            scales = np.zeros((k_blocks, cols), dtype=fp32weight.dtype)
            quantize_qdq_matmul_4bits(
                packed, fp32weight, scales, zero_point, block_size, cols, rows, self.config.is_symmetric
            )

        return (packed, scales, zero_point)

    def quantize_matmul(self, node: NodeProto, graph_stack: list[GraphProto]) -> list[NodeProto]:
        """
        Quantize weight B of MatMul node to int4 or int8.
        Supports 2D constant matrix, and N-D constant matrices whose leading dimensions are all 1
        (which are squeezed to 2D before quantization). Axis 0 blockwise quantization only.
        """
        bits = self.config.bits
        if bits == 8:
            qtype = TensorProto.INT8 if self.config.is_symmetric else TensorProto.UINT8
        else:
            qtype = TensorProto.INT4 if self.config.is_symmetric else TensorProto.UINT4
        input_b = node.input[1]
        b_tensor, b_graph = get_initializer(input_b, graph_stack)
        if b_tensor is None:
            logger.info("MatMul doesn't have const weight. Skip to quantize")
            return [node]  # only care about constant weight

        b_ndarray = ir.from_proto(b_tensor).numpy()
        b_original_shape = b_ndarray.shape
        if len(b_original_shape) != 2:
            if len(b_original_shape) < 2:
                logger.info("MatMul weight has fewer than 2 dimensions. Skip to quantize.")
                return [node]
            leading = b_original_shape[:-2]
            if any(d != 1 for d in leading):
                logger.info(
                    "MatMul weight has non-unit batch dims %s; N-D batched quantization not supported. "
                    "Skip to quantize.",
                    list(leading),
                )
                return [node]
            # Squeeze all unit leading dims to obtain a 2-D weight [K, N]
            b_ndarray = b_ndarray.reshape(b_original_shape[-2], b_original_shape[-1])
        else:
            b_original_shape = None  # already 2-D, no reshape needed

        bfloat16 = b_ndarray.dtype == "bfloat16"
        if bfloat16:
            b_ndarray = b_ndarray.astype(np.float32)

        packed, scales, zero_points = self.qbits_block_quant(b_ndarray)
        if bfloat16:
            scales = scales.astype(ml_dtypes.bfloat16)

        if self.config.quant_format == QuantFormat.QOperator:
            b_quant = ir.serde.serialize_tensor(ir.Tensor(packed, name=b_tensor.name + f"_Q{bits}"))
            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_scales"))
        else:
            b_quant = onnx.helper.make_tensor(
                b_tensor.name + f"_DQ_Q{bits}", qtype, b_ndarray.shape, packed.tobytes(), True
            )
            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_DQ_scales"))

        # if QDQ, CW and SYM enabled, optimize for Intel NPU, tranpose the weight to NHWC format will increase performance
        qdq_opt_for_intel_npu_enabled = (
            self.config.quant_format == QuantFormat.QDQ
            and self.config.channel_wised_quantize
            and self.config.is_symmetric
        )
        if qdq_opt_for_intel_npu_enabled:
            rows, cols = b_ndarray.shape
            packed = transpose_packed_int4_matrix(packed, rows, cols)
            scales = scales.reshape((cols, 1))  # (cols, 1)
            b_quant = onnx.helper.make_tensor(
                b_tensor.name + f"_DQ_Q{bits}", qtype, [cols, rows], packed.tobytes(), True
            )
            scales_tensor = ir.serde.serialize_tensor(ir.Tensor(scales, name=b_tensor.name + "_DQ_scales"))

        for input in b_graph.input:
            if input.name == input_b:
                b_graph.input.remove(input)
                break

        b_graph.initializer.extend([b_quant, scales_tensor])

        output_nodes = []

        if self.config.quant_format == QuantFormat.QOperator:
            input_names = [node.input[0], b_quant.name, scales_tensor.name]
            if not self.config.is_symmetric:
                zp_tensor = onnx.numpy_helper.from_array(zero_points, b_tensor.name + "_zero_points")
                input_names.append(zp_tensor.name)
                b_graph.initializer.extend([zp_tensor])
            kwargs = {}
            rows, cols = b_ndarray.shape
            kwargs["K"] = rows
            kwargs["N"] = cols
            kwargs["bits"] = bits
            kwargs["block_size"] = self.config.block_size

            # Do not output accuracy_level if it is 0 since the attribute is optional and is not supported by most EPs.
            if self.config.accuracy_level:
                kwargs["accuracy_level"] = self.config.accuracy_level

            # When rank(A) >= rank_b_orig, the reshape would be a no-op since MatMulNBits
            # already produces the correct output shape.
            a_static_rank = _get_static_rank(node.input[0], graph_stack)
            rank_b_orig = len(b_original_shape) if b_original_shape is not None else 0
            needs_reshape = b_original_shape is not None and (a_static_rank is None or a_static_rank < rank_b_orig)
            qop_output = node.output[0] if not needs_reshape else node.output[0] + "_pre_reshape"
            matmul_qbit_node = onnx.helper.make_node(
                "MatMulNBits",
                inputs=input_names,
                outputs=[qop_output],
                name=node.name + f"_Q{bits}" if node.name else "",
                domain="com.microsoft",
                **kwargs,
            )

            output_nodes.append(matmul_qbit_node)
            if needs_reshape:
                output_nodes.extend(
                    _build_nbits_output_reshape(
                        a_input_name=node.input[0],
                        b_original_shape=b_original_shape,
                        target_graph=b_graph,
                        name_prefix=(node.name + f"_Q{bits}") if node.name else (b_tensor.name + f"_Q{bits}"),
                        pre_reshape_output=qop_output,
                        final_output=node.output[0],
                    )
                )
        else:
            dq_input_names = [b_quant.name, scales_tensor.name]
            dq_output_names = [b_quant.name + "_output"]
            tp_input_names = [dq_output_names[0]]
            tp_output_names = [dq_output_names[0] + "_transposed"]
            matmul_input_names = [
                node.input[0],
                tp_output_names[0] if qdq_opt_for_intel_npu_enabled else dq_output_names[0],
            ]
            # When rank(A) >= rank_b_orig, the reshape would be a no-op since MatMul
            # already produces the correct output shape.
            a_static_rank = _get_static_rank(node.input[0], graph_stack)
            rank_b_orig = len(b_original_shape) if b_original_shape is not None else 0
            needs_reshape = b_original_shape is not None and (a_static_rank is None or a_static_rank < rank_b_orig)
            qdq_matmul_out = node.output[0] if not needs_reshape else node.output[0] + "_pre_reshape"
            matmul_output_names = [qdq_matmul_out]
            if not self.config.is_symmetric:
                zp_tensor = onnx.helper.make_tensor(
                    b_tensor.name + "_DQ_zero_points", qtype, scales.shape, zero_points.tobytes(), True
                )
                dq_input_names.append(zp_tensor.name)
                b_graph.initializer.extend([zp_tensor])
            rows, cols = b_ndarray.shape
            dq_kwargs = {
                "axis": 1 if qdq_opt_for_intel_npu_enabled else 0,
                "block_size": rows if self.config.channel_wised_quantize else self.config.block_size,
            }
            dq_node = onnx.helper.make_node(
                "DequantizeLinear",
                inputs=dq_input_names,
                outputs=dq_output_names,
                name=node.name + f"_DQ_Q{bits}" if node.name else "",
                **dq_kwargs,
            )
            matmul_node = onnx.helper.make_node(
                "MatMul",
                inputs=matmul_input_names,
                outputs=matmul_output_names,
                name=node.name + f"_matmul_Q{bits}" if node.name else "",
            )
            if qdq_opt_for_intel_npu_enabled:
                tp_node = onnx.helper.make_node(
                    "Transpose",
                    inputs=tp_input_names,
                    outputs=tp_output_names,
                    perm=[1, 0],
                )
                output_nodes.extend([dq_node, tp_node, matmul_node])
            else:
                output_nodes.extend([dq_node, matmul_node])
            if needs_reshape:
                output_nodes.extend(
                    _build_nbits_output_reshape(
                        a_input_name=node.input[0],
                        b_original_shape=b_original_shape,
                        target_graph=b_graph,
                        name_prefix=(node.name + f"_DQ_Q{bits}") if node.name else (b_tensor.name + f"_DQ_Q{bits}"),
                        pre_reshape_output=qdq_matmul_out,
                        final_output=node.output[0],
                    )
                )

        return output_nodes

    @staticmethod
    def quant_slice_symmetric(data: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        max_val = np.max(data, axis=1, keepdims=True)
        min_val = np.min(data, axis=1, keepdims=True)
        abs_max = np.where(np.abs(max_val) > np.abs(min_val), max_val, min_val)

        scale = abs_max / -8.0  # if max == min, max may be clipped
        quantized_slice = np.where(scale == 0, 0, data / scale).round().clip(-8, 7).astype(np.int8)

        return quantized_slice, scale

    @staticmethod
    def quant_slice_asymmetric(data: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        min_val = np.minimum(data.min(axis=1, keepdims=True), 0)
        max_val = np.maximum(data.max(axis=1, keepdims=True), 0)

        scale = (max_val - min_val) / 15.0
        zero_point = np.where(scale == 0, 8, -min_val / scale).round().clip(0, 15).astype(np.uint8)
        quantized_slice = np.where(scale == 0, 8, data / scale + zero_point).round().clip(0, 15).astype(np.uint8)

        return quantized_slice, scale, zero_point

    @staticmethod
    def pack_int8_to_int4(data: np.ndarray) -> np.ndarray:
        """Pack int8 data to int4 and store in uint8 ndarray."""
        data_flat = data.reshape(-1)
        if len(data_flat) % 2 != 0:
            data_flat = np.append(data_flat, 0)
        quant_data_int4 = (data_flat[::2] & 0xF) | ((data_flat[1::2] & 0xF) << 4)

        return quant_data_int4.astype("uint8")

    @staticmethod
    def quantize_ndarray(
        data: np.ndarray,
        quantize_axis: int,
        block_size: int,
        is_symmetric: bool,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
        """Quantize ndarray data to int4 using numpy, return (quantized data, scales, zero points)."""
        # Get the shape of the matrix
        m = 1  # dimension of the matrix before the quantize axis
        k = data.shape[quantize_axis]  # dimension of the matrix along the quantize axis
        n = 1  # dimension of the matrix after the quantize axis
        for i, dim in enumerate(data.shape):
            if i < quantize_axis:
                m *= dim
            elif i > quantize_axis:
                n *= dim

        k_blocks = (k + block_size - 1) // block_size
        scales_shape = list(data.shape)
        scales_shape[quantize_axis] = k_blocks

        data_reshape = data.reshape((m, k, n))
        scales = np.zeros((m, k_blocks, n), dtype=data.dtype)
        if is_symmetric:
            quant_data_int8 = np.zeros((m, k, n), dtype="int8")
        else:
            quant_data_int8 = np.zeros((m, k, n), dtype="uint8")
            zero_point_int8 = np.zeros((m, k_blocks, n), dtype="uint8")

        # slice and quantize
        for i in range(0, k, block_size):
            end_idx = min(i + block_size, k)
            slice = data_reshape[:, i:end_idx, :]

            if is_symmetric:
                quantized_slice_int8, scale_slice = DefaultWeightOnlyQuantizer.quant_slice_symmetric(slice)
            else:
                quantized_slice_int8, scale_slice, zero_point_slice_int8 = (
                    DefaultWeightOnlyQuantizer.quant_slice_asymmetric(slice)
                )

            quant_data_int8[:, i:end_idx, :] = quantized_slice_int8
            j = i // block_size
            scales[:, j : (j + 1), :] = scale_slice
            if not is_symmetric:
                zero_point_int8[:, j : (j + 1), :] = zero_point_slice_int8

        # pack int8 to int4
        quant_data_int4 = DefaultWeightOnlyQuantizer.pack_int8_to_int4(quant_data_int8)
        zero_point_int4 = None
        if not is_symmetric:
            zero_point_int4 = DefaultWeightOnlyQuantizer.pack_int8_to_int4(zero_point_int8)
        scales = scales.reshape(scales_shape)
        return quant_data_int4, scales, zero_point_int4

    def quantize_gather(self, node: NodeProto, graph_stack: list[GraphProto]) -> list[NodeProto]:
        """Quantize weight data of Gather node to int4."""
        assert self.config.quant_format == QuantFormat.QOperator, "Gather only supports QOperator format currently."

        qtype = TensorProto.INT4 if self.config.is_symmetric else TensorProto.UINT4
        data_arg = node.input[0]
        data_tensorproto, data_graphproto = get_initializer(data_arg, graph_stack)
        if data_tensorproto is None:
            logger.info("Gather doesn't have const weight. Skip quantization.")
            return [node]  # only care about constant weight

        data_ndarray = onnx.numpy_helper.to_array(data_tensorproto)
        data_rank = len(data_ndarray.shape)
        quantize_axis = self.config.quant_axes.get("Gather", 1)
        block_size = self.config.block_size

        assert quantize_axis < data_rank and quantize_axis >= -data_rank, "Invalid quantize axis for Gather node."
        assert block_size >= 16 and ((block_size - 1) & block_size == 0), "Invalid block size for Gather node."

        quantize_axis = (quantize_axis + data_rank) % data_rank
        quantized_data, scales, zero_points = self.quantize_ndarray(
            data_ndarray, quantize_axis, block_size, self.config.is_symmetric
        )

        for input in data_graphproto.input:
            if input.name == data_arg:
                data_graphproto.input.remove(input)
                break

        quantized_data_tensorproto = onnx.helper.make_tensor(
            data_tensorproto.name + "_Q4", qtype, data_ndarray.shape, quantized_data.tobytes(), True
        )
        scales_tensorproto = onnx.numpy_helper.from_array(scales, data_tensorproto.name + "_scales")
        input_names = [quantized_data_tensorproto.name, node.input[1], scales_tensorproto.name]
        data_graphproto.initializer.extend([quantized_data_tensorproto, scales_tensorproto])
        if not self.config.is_symmetric:
            zp_tensorproto = onnx.helper.make_tensor(
                data_tensorproto.name + "_zero_points", qtype, scales.shape, zero_points.tobytes(), True
            )
            input_names.append(zp_tensorproto.name)
            data_graphproto.initializer.extend([zp_tensorproto])

        try:
            gather_axis = onnx.helper.get_node_attr_value(node, "axis")
        except ValueError:
            gather_axis = 0

        kwargs = {
            "gather_axis": gather_axis,
            "quantize_axis": quantize_axis,
            "block_size": block_size,
        }

        gather_q4_node = onnx.helper.make_node(
            "GatherBlockQuantized",
            inputs=input_names,
            outputs=[node.output[0]],
            name=node.name + "_Q4" if node.name else "",
            domain="com.microsoft",
            **kwargs,
        )

        return [gather_q4_node]

    def quantize(self, node: NodeProto, graph_stack: list[GraphProto]) -> list[NodeProto]:
        """
        Target node:        QOperator node:            QDQ nodes:
        MatMul              MatMulNBits                DeQuantizeLinear -> MatMul
        Gather              GatherBlockQuantized       Gather, Gather, Gather (optional) -> DequantizeLinear
        If the node is target node with fp32 or fp16 const weight, quantize the weight to int4 and
        return the new nodes.
        If QOperator format, return the corresponding QOperator nodes.
        If QDQ format, return the corresdponging QDQ nodes.
        Gather (quantized data) + Gather (scales) + Gather (optional, zero points) -> DequantizeLinear is
        not supported yet because Gather does not support int4 data.
        """
        logger.info(f"start to quantize {node.name} ...")

        bits = self.config.bits
        if node.op_type == "MatMul":
            if bits == 8 and self.config.quant_format == QuantFormat.QDQ:
                logger.error("MatMul only supports QOperator format for 8 bits quantization.")
                return [node]
            results = self.quantize_matmul(node, graph_stack)
        elif node.op_type == "Gather":
            if self.config.bits != 4:
                logger.error("Gather only supports 4 bits quantization.")
                return [node]

            results = self.quantize_gather(node, graph_stack)
        else:
            logger.error(f"Unsupported operator {node.op_type} for weight only quantization. Skip quantization.")
            return [node]

        logger.info(f"complete quantization of {node.name} with {self.config.bits} bits ...")
        return results


class NVAWQWeightOnlyQuantizer:
    def __init__(
        self,
        config: NVAWQWeightOnlyQuantConfig,
    ):
        self.config = config

    def quantize_awq(self, model: ModelProto | str) -> ModelProto:
        """
        Perform nvidia_awq quantization using ModelOpt's int4 quantize function.

        Args:
            model (ModelProto): The ONNX model to quantize.

        Returns:
            ModelProto: The quantized ONNX model.
        """
        try:
            from modelopt.onnx.quantization.int4 import quantize as quantize_int4  # noqa: PLC0415
        except ImportError:
            print(
                "Please ensure that the 'modelopt' package is installed. Please install it using pip install nvidia_modelopt."
            )
            raise ImportError(
                "modelopt is not installed. Please install it using pip install nvidia_modelopt. Exiting."
            ) from None

        logger.info("Starting nvidia_awq quantization...")

        # Prepare calibration inputs
        calib_inputs = self.config.calibration_data_reader

        # Perform quantization using ModelOpt's int4 quantize function
        quantized_model = quantize_int4(
            model,
            calibration_method=self.config.calibration_method,
            calibration_data_reader=calib_inputs,
        )

        logger.info("Completed nvidia_awq quantization.")
        return quantized_model


class MatMulNBitsQuantizer:
    """
    Target node:        QOperator node:            QDQ nodes:
    MatMul              MatMulNBits                DeQuantizeLinear -> MatMul
    Gather              GatherBlockQuantized       Gather, Gather, Gather (optional) -> DequantizeLinear

    Perform 2/4/8 bits quantization of constant weights for target nodes.
    If algo_config.quant_format is QOperator:
      - nodes are replaced by the corresponding QOperator nodes.
      - quantized weights are stored in the contrib ops.
    If algo_config.quant_format is QDQ:
      - the quantized weight is stored in a standard onnx node. For MatMul, it is DequantizeLinear. For Gather,
        it is the three Gathers, one for quantized data, one for scales and one for optional zero points.
      - The nodes are replaced by the corresponding QDQ nodes.
      - currently Gather is not supported in QDQ because Gather does not support int4 yet.
    Note:
      - for quantized gather, the memory usage of "DequantizeLinear + Gather" is the same as the original Gather
        during runtime. Therefor it is not recommended.
      - when a node is in nodes_to_exclude, and the node configuration in algo_config.customized_weight_config will be ignored.
    """

    def __init__(
        self,
        model: ModelProto | str,
        bits: int = 4,  # default to 4bit
        block_size: int = 128,
        is_symmetric: bool = False,
        accuracy_level: int | None = None,
        nodes_to_exclude=None,
        nodes_to_include: list[str] | None = None,
        quant_format=QuantFormat.QOperator,
        op_types_to_quantize: tuple[str, ...] | None = None,
        quant_axes: tuple[tuple[str, int], ...] | None = None,
        channel_wised_quantize: bool = False,
        algo_config: WeightOnlyQuantConfig | None = None,
    ):
        if nodes_to_exclude is None:
            nodes_to_exclude = []
        self.model = ONNXModel(onnx.load(model)) if isinstance(model, str) else ONNXModel(model)
        self.model_path = model if isinstance(model, str) else None
        self.bits = bits
        self.block_size = block_size
        self.is_symmetric = is_symmetric
        self.accuracy_level = accuracy_level
        self.nodes_to_exclude = set(nodes_to_exclude)
        self.nodes_to_include = set(nodes_to_include) if nodes_to_include else None
        self.node_quantizer = None

        if algo_config is None:
            algo_config = DefaultWeightOnlyQuantConfig(
                block_size=block_size,
                is_symmetric=is_symmetric,
                accuracy_level=accuracy_level,
                quant_format=quant_format,
                op_types_to_quantize=op_types_to_quantize,
                quant_axes=quant_axes,
                bits=bits,
                channel_wised_quantize=channel_wised_quantize,
            )

        self.algo_config = algo_config
        if hasattr(self.algo_config, "bits"):
            assert self.algo_config.bits in [2, 4, 8], "Only support 2, 4 or 8 bits quantization"

        if algo_config.algorithm == "HQQ":
            self.node_quantizer = HQQWeightOnlyQuantizer(self.algo_config)
        elif algo_config.algorithm == "DEFAULT":
            self.node_quantizer = DefaultWeightOnlyQuantizer(self.algo_config)
        elif algo_config.algorithm == "nvidia_awq":
            self.node_quantizer = NVAWQWeightOnlyQuantizer(self.algo_config)

    def _process_subgraph(self, graph_stack: list[GraphProto]):
        new_nodes = []
        graph = graph_stack[-1]

        for node in graph.node:
            graph_attrs = [
                attr
                for attr in node.attribute
                if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS
            ]
            if graph_attrs:
                kwargs = {}
                for attr in node.attribute:
                    if attr.type == onnx.AttributeProto.GRAPH:
                        # recursive call to take care of sub-graph
                        graph_stack.append(attr.g)
                        kv = {attr.name: self._process_subgraph(graph_stack)}
                    elif attr.type == onnx.AttributeProto.GRAPHS:
                        value = []
                        for subgraph in attr.graphs:
                            # recursive call to take care of sub-graph
                            graph_stack.append(subgraph)
                            value.extend([self._process_subgraph(graph_stack)])
                        kv = {attr.name: value}
                    else:
                        kv = attribute_to_kwarg(attr)
                    kwargs.update(kv)
                node = onnx.helper.make_node(  # noqa: PLW2901
                    node.op_type, node.input, node.output, name=node.name, **kwargs
                )
            out_nodes = []
            if node.name in self.nodes_to_exclude:
                logger.info(f"exclude to quantize {node.name} as specified by nodes_to_exclude...")
                out_nodes = [node]
            elif (self.nodes_to_include and node.name in self.nodes_to_include) or (
                node.op_type in self.algo_config.op_types_to_quantize
            ):
                out_nodes = self.node_quantizer.quantize(node, graph_stack)
            else:
                logger.info(f"skip to quantize {node.name} ...")
                out_nodes = [node]
            new_nodes.extend(out_nodes)

        graph.ClearField("node")
        graph.node.extend(new_nodes)
        graph_stack.pop()
        return graph

    def _generate_q4_node_config(self):
        """Generate weight only quant configuration for nodes."""
        q4_node_config = {}
        for node in self.model.model.graph.node:
            if node.op_type in ["MatMul"]:
                if not all(self.model.get_initializer(i) is None for i in node.input):
                    template_config_q4 = {
                        "bits": 4,
                        "group_size": self.block_size,
                        "scheme": "sym" if self.is_symmetric else "asym",
                    }
                    if (
                        self.algo_config.customized_weight_config
                        and node.name in self.algo_config.customized_weight_config
                    ):
                        for key, value in self.algo_config.customized_weight_config[node.name].items():
                            if key in template_config_q4:
                                template_config_q4[key] = value
                    q4_node_config[node.name] = template_config_q4
        return q4_node_config

    def int4_quant_algo(self):
        """4b quantize a model with RTN or GPTQ algorithm. Please refer to
        https://github.com/intel/neural-compressor/blob/master/docs/source/quantization_weight_only.md
        for more details on weight only quantization using Intel® Neural Compressor.
        """

        def inc_dataloader():
            data_reader = copy.deepcopy(self.algo_config.calibration_data_reader)
            for data in data_reader:
                yield data, None

        kwargs = {}
        if self.accuracy_level is not None:
            kwargs["accuracy_level"] = self.accuracy_level
        weight_only_node_config = self._generate_q4_node_config()

        algorithm = self.algo_config.algorithm
        logger.info(f"start to quantize model with {algorithm} algorithm...")
        if algorithm in ["RTN", "k_quant"]:
            kwargs["ratios"] = self.algo_config.ratios
            kwargs["algorithm"] = algorithm

            """
            We uses fp32 to represent the node that skip quantization, it does not mean this node is fp32 type though.
            """
            for n in self.nodes_to_exclude:
                weight_only_node_config[n] = "fp32"

            self.model = rtn_quantize(
                model=self.model_path if self.model_path is not None else self.model.model,
                weight_config=weight_only_node_config,
                **kwargs,
            )
        elif algorithm == "GPTQ":
            kwargs["percdamp"] = self.algo_config.percdamp
            kwargs["blocksize"] = self.algo_config.block_size
            kwargs["actorder"] = self.algo_config.actorder
            kwargs["mse"] = self.algo_config.mse
            kwargs["perchannel"] = self.algo_config.perchannel
            kwargs["n_samples"] = -1
            dataloader = inc_dataloader()

            self.model = gptq_quantize(
                model=self.model_path if self.model_path is not None else self.model.model,
                weight_config=weight_only_node_config,
                dataloader=dataloader,
                **kwargs,
            )
        logger.info(f"complete quantization of model with {algorithm} algorithm.")

    def process(self):
        if self.algo_config.algorithm in ["HQQ", "DEFAULT"]:
            # use a stack to keep track of sub-graphs
            graph_stack = [self.model.graph()]

            # Update domain opset
            if self.algo_config.quant_format == QuantFormat.QOperator:
                self.model.set_opset_import("com.microsoft", 1)

            if self.algo_config.quant_format == QuantFormat.QDQ or "Gather" in self.algo_config.op_types_to_quantize:
                opset_import = self.model.opset_import()
                for opset in opset_import:
                    if opset.domain in [None, "ai.onnx", ""] and opset.version < 21:
                        logger.warning(
                            "The opset of the input model is under 21 and doesn't support int4 data type. "
                            "Force to update it to opset 21, but the generated model may not be a valid model."
                        )
                        self.model.set_opset_import(opset.domain, 21)

            self._process_subgraph(graph_stack)
            self.model.clean_initializers()
        elif self.algo_config.algorithm == "nvidia_awq":
            # Handle nvidia_awq quantization
            logger.info("Processing nvidia_awq quantization...")
            self.model = self.node_quantizer.quantize_awq(
                self.model.model if self.model_path is None else self.model_path
            )
            logger.info("Completed nvidia_awq quantization.")
            self.model = ONNXModel(self.model)  # Ensure the model is wrapped back into ONNXModel
            self.model.clean_initializers()
        else:
            # RTN or GPTQ weight-only quantize algorithm
            self.int4_quant_algo()


def ort_convert_str_to_bool(value):
    return value.lower() in ("true", "1")


# Custom function to parse str:int pairs
def parse_key_value_pair(s):
    key, value = s.split(":")
    return key, int(value)


def parse_args():
    parser = argparse.ArgumentParser(
        description="""Blockwise int4 quantization for MatMul 2D weight matrices.

A weight matrix is partitioned into into blocks, where each block is a
continguous subset inside each column. Each block is quantized into a
set of 4b integers with a scaling factor and an optional offset.
"""
    )

    parser.add_argument("--input_model", required=True, help="Path to the input model file")
    parser.add_argument("--output_model", required=True, help="Path to the output model file")
    parser.add_argument("--block_size", required=False, default=32, type=int, help="Block size for quantization")
    parser.add_argument(
        "--quant_method",
        default="default",
        type=str,
        choices=["default", "hqq", "rtn", "k_quant", "gptq", "nvidia_awq"],
        help="the algorithm used to quantize weight, \nrtn and gptq leverage Intel® Neural Compressor",
    )
    parser.add_argument("--bits", default=4, type=int, help="the target bits to represent weight")
    parser.add_argument(
        "--symmetric",
        required=False,
        default=True,
        const=True,
        nargs="?",
        type=ort_convert_str_to_bool,
        choices=[True, False],
        help="Indicate whether to quantize the model symmetrically, symmetric is not supported by hqq",
    )
    parser.add_argument(
        "--accuracy_level",
        required=False,
        type=int,
        help="Accuracy level of the 4-bit quantized MatMul computation. "
        "Refer to the MatMulNBits contrib op's 'accuracy_level' attribute for details "
        "(https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftmatmulnbits).",
    )
    parser.add_argument("-v", "--verbose", required=False, action="store_true")
    parser.set_defaults(verbose=False)
    parser.add_argument(
        "--nodes_to_exclude",
        nargs="+",
        type=str,
        required=False,
        default=[],
        help="Specify the nodes to be excluded from quantization with node names",
    )
    parser.add_argument(
        "--nodes_to_include",
        nargs="+",
        type=str,
        required=False,
        help="Specify the specific nodes to be included from quantization with node names",
    )
    parser.add_argument(
        "--quant_format",
        default="QOperator",
        type=str,
        choices=["QOperator", "QDQ"],
        help="QuantFormat {QOperator, QDQ}"
        "QOperator format quantizes the model with quantized operators directly."
        "QDQ format quantize the model by inserting DeQuantizeLinear before the MatMul.",
    )
    parser.add_argument(
        "--op_types_to_quantize",
        type=str,
        nargs="+",
        choices=["MatMul", "Gather"],
        help="op_types_to_quantize {MatMul, Gather}. Operators to quantize. Default is MatMul.",
    )
    parser.add_argument(
        "--quant_axes",
        type=parse_key_value_pair,
        nargs="+",
        required=False,
        help="Key-value pairs in op_type:axis_to_quantize separated by space."
        "Specify the axis to quantize for an op. Default {MatMul:0, Gather:1}"
        "Example: --quant_axes MatMul:0 Gather:1",
    )
    # Group arguments specific to nvidia_awq
    nv_awq_config = parser.add_argument_group("nvidia_awq", "Arguments specific to nvidia_awq quantization")
    nv_awq_config.add_argument(
        "--calib_dataset_name",
        type=str,
        default="cnn",
        help="Name of the calibration dataset for nvidia_awq.",
    )
    nv_awq_config.add_argument(
        "--tokenizer_dir",
        type=str,
        required=False,
        help="Path of the tokenizer dir.",
    )
    nv_awq_config.add_argument(
        "--calibration_method",
        type=str,
        required=False,
        choices=["awq", "awq_clip"],
        help="Support two options, awq implementation and weight clipping.",
    )
    nv_awq_config.add_argument(
        "--cache_dir",
        type=str,
        default="./cache",
        help="Cache directory for calibration data.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    if args.verbose:
        logger.setLevel(logging.DEBUG)

    input_model_path = args.input_model
    output_model_path = args.output_model
    quant_format = QuantFormat[args.quant_format]
    op_types_to_quantize = tuple(args.op_types_to_quantize) if args.op_types_to_quantize else ("MatMul",)
    quant_axes = tuple(args.quant_axes) if args.quant_axes else None

    if os.path.exists(output_model_path):
        logger.error(f"file {output_model_path} already exists")
        raise Exception(f"file {output_model_path} already exists")

    if args.symmetric and args.quant_method == "hqq":
        logger.warning("symmetric is not supportted by hqq, will force to symmetric=False")
        args.symmetric = False

    model = onnx.load(input_model_path)
    if args.quant_method == "hqq":
        quant_config = HQQWeightOnlyQuantConfig(
            block_size=args.block_size, bits=args.bits, op_types_to_quantize=op_types_to_quantize, quant_axes=quant_axes
        )
    elif args.quant_method == "default":
        quant_config = DefaultWeightOnlyQuantConfig(
            block_size=args.block_size,
            is_symmetric=args.symmetric,
            accuracy_level=args.accuracy_level,
            quant_format=quant_format,
            op_types_to_quantize=op_types_to_quantize,
            quant_axes=quant_axes,
            bits=args.bits,
        )
    elif args.quant_method == "rtn":
        quant_config = RTNWeightOnlyQuantConfig(op_types_to_quantize=op_types_to_quantize)
    elif args.quant_method == "k_quant":
        quant_config = KQuantWeightOnlyQuantConfig(op_types_to_quantize=op_types_to_quantize)
    elif args.quant_method == "gptq":
        quant_config = GPTQWeightOnlyQuantConfig(block_size=args.block_size, op_types_to_quantize=op_types_to_quantize)
    elif args.quant_method == "nvidia_awq":
        if quant_format == QuantFormat.QOperator:
            logger.warning("QOperator is not applicable to nvidia_awq. overriding the value to QDQ")
            quant_format = QuantFormat.QDQ

        model = input_model_path
        if args.calibration_method is not None:
            if args.calibration_method == "awq":
                calibration_method = "awq_lite"
            else:
                calibration_method = "awq_clip"
        else:
            calibration_method = "awq_lite"

        quant_config = NVAWQWeightOnlyQuantConfig(
            dataset_name=args.calib_dataset_name,
            tokenizer_dir=args.tokenizer_dir,
            cache_dir=args.cache_dir,
            calibration_method=calibration_method,
        )
    else:
        raise ValueError(f"Unsupported quantization method: {args.quant_method}")

    quant = MatMulNBitsQuantizer(
        model=model,
        bits=args.bits,
        accuracy_level=args.accuracy_level,
        nodes_to_exclude=args.nodes_to_exclude,
        nodes_to_include=args.nodes_to_include,
        algo_config=quant_config,
    )
    quant.process()
    quant.model.save_model_to_file(output_model_path, True)