# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.  See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# Modifications Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
# Portions of this file consist of AI generated content.
"""
A set of Python classes to unpack the quantized weights and repack them in ONNX Runtime's
standard format.

The goal is for `QuantModel` to repack the quantized weights into a standard format
so that the original Hugging Face --> ONNX code can store the quantized weights as
ONNX Runtime's format no matter where the quantized weights actually come from.
"""

import os
import re

import torch
from safetensors.torch import load_file


def normalize_vlm_weight_name(name):
    """Normalize a checkpoint tensor key for VLM/Quark conventions.

    Returns None if the tensor should be skipped (vision-tower weights), or
    the normalized key string otherwise.
    """
    # Skip vision tower weights in VLM checkpoints
    if name.startswith(("model.visual.", "model.vision.", "visual.")):
        return None
    # Normalize common VLM prefix so existing LLM regex + parsing keeps working
    if name.startswith("model.language_model."):
        name = "model." + name[len("model.language_model."):]
    # Normalize Quark weight_quantizer.* naming to flat weight_* naming
    name = name.replace(".weight_quantizer.scale", ".weight_scale")
    name = name.replace(".weight_quantizer.zero_point", ".weight_zero_point")
    return name


class QuantizedTensorModule:
    def __init__(self):
        self.qweight = None
        self.scales = None
        self.qzeros = None
        self.g_idx = None
        self.bias = None

        self.in_features = 0
        self.out_features = 0
        self.bits = None
        self._group_size = None

    @property
    def group_size(self):
        return self._group_size if self._group_size != -1 else self.in_features

    @group_size.setter
    def group_size(self, value):
        self._group_size = value

    def __str__(self):
        qweight = f"qweight = {self.qweight.shape}, {self.qweight}\n"
        scales = f"scales = {self.scales.shape}, {self.scales}\n"
        qzeros = "" if self.qzeros is None else f"qzeros = {self.qzeros.shape}, {self.qzeros}\n"
        g_idx = "" if self.g_idx is None else f"g_idx = {self.g_idx.shape}, {self.g_idx}\n"

        in_feats = f"in_features = {self.in_features}, "
        out_feats = f"out_features = {self.out_features}, "
        bits = f"bits = {self.bits}, "
        group_size = f"group_size = {self.group_size}, "

        return qweight + qzeros + scales + g_idx + in_feats + out_feats + bits + group_size


class TensorModule:
    def __init__(self):
        self.weight = None
        self.bias = None


class QuantizedAttention:
    def __init__(self):
        self.q_proj = QuantizedTensorModule()
        self.k_proj = QuantizedTensorModule()
        self.v_proj = QuantizedTensorModule()
        self.o_proj = QuantizedTensorModule()
        self.rotary_emb = TensorModule()
        self.k_norm = TensorModule()
        self.q_norm = TensorModule()
        self.sinks = TensorModule()


class QuantizedExpert:
    """Represents a single expert in MoE with quantized weights."""
    def __init__(self, expert_id: int):
        self.expert_id = expert_id
        self.gate_proj = QuantizedTensorModule()
        self.up_proj = QuantizedTensorModule()
        self.gate_up_proj = QuantizedTensorModule()
        self.down_proj = QuantizedTensorModule()


class QuantizedExperts:
    """Container for all experts in a MoE layer."""
    def __init__(self):
        """Pre-processed experts attributes"""
        self._experts = {}
        """QMoE packed attributes"""
        self.fc1_weights = None
        self.fc1_scales = None
        self.fc1_zero_points = None
        self.fc2_weights = None
        self.fc2_scales = None
        self.fc2_zero_points = None

    def add_expert(self, expert_id: int) -> QuantizedExpert:
        """Add a new expert and return it."""
        if expert_id not in self._experts:
            self._experts[expert_id] = QuantizedExpert(expert_id)
        return self._experts[expert_id]

    def get_expert(self, expert_id: int) -> QuantizedExpert:
        """Get an expert by ID, creating if it doesn't exist."""
        return self.add_expert(expert_id)

    def items(self):
        """Get (expert_id, expert) pairs."""
        return self._experts.items()

    def values(self):
        """Get all experts."""
        return self._experts.values()

    def keys(self):
        """Get all expert IDs."""
        return self._experts.keys()

    def __getitem__(self, expert_id: int) -> QuantizedExpert:
        """Get expert by ID."""
        return self.get_expert(expert_id)

    @property
    def num_experts(self) -> int:
        """Number of experts."""
        return len(self._experts)

    def set_weight_data(self, expert_id: int, proj_type: str, param_type: str, tensor, bits: int, group_size: int):
        """Set weight data for a specific expert projection.

        Args:
            expert_id: Expert ID (0, 1, 2, ...)
            proj_type: Projection type ('gate_proj', 'up_proj', 'down_proj')
            param_type: Parameter type ('weight', 'weight_zero_point', 'scales', etc.)
            tensor: The tensor data
            bits: Quantization bits
            group_size: Quantization group size
        """
        expert = self.get_expert(expert_id)
        proj_module = getattr(expert, proj_type)

        proj_module.bits = bits
        proj_module.group_size = group_size

        # Map parameter names
        param_mapping = {
            'weight': 'qweight',
            'weight_zero_point': 'qzeros',
            'weight_scale': 'scales',
            'scales': 'scales',
            'qweight': 'qweight',
            'qzeros': 'qzeros',
            'bias': 'bias',
            'g_idx': 'g_idx'
        }

        attr_name = param_mapping.get(param_type, param_type)
        setattr(proj_module, attr_name, tensor)

    def __str__(self):
        """String representation of all experts in the MoE layer."""
        if not self._experts:
            return "QuantizedExperts(num_experts=0)"

        lines = [f"QuantizedExperts(num_experts={self.num_experts})"]
        for expert_id, expert in sorted(self._experts.items()):
            lines.append(f"  QuantizedExperts {expert_id}:")
            lines.append(f"  gate_proj: {expert.gate_proj}")
            lines.append(f"  up_proj: {expert.up_proj}")
            lines.append(f"  down_proj: {expert.down_proj}")

        return "\n".join(lines)


class QuantizedMLP:
    def __init__(self):
        self.gate_proj = QuantizedTensorModule()
        self.up_proj = QuantizedTensorModule()
        self.down_proj = QuantizedTensorModule()
        self.fc1 = QuantizedTensorModule()
        self.fc2 = QuantizedTensorModule()
        # MoE
        self.experts = QuantizedExperts()
        self.router = TensorModule()


class QuantizedDecoderLayer:
    def __init__(self, layer_id):
        self.layer_id = layer_id
        self.input_layernorm = TensorModule()
        self.self_attn = QuantizedAttention()
        self.post_attention_layernorm = TensorModule()
        self.pre_feedforward_layernorm = TensorModule()
        self.post_feedforward_layernorm = TensorModule()
        self.mlp = QuantizedMLP()

    def is_empty(self):
        return self.input_layernorm.weight is None


class QuantizedModel:
    def __init__(
        self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers, load_weights=True
    ):
        self.quant_type = quant_type
        self.embedding = TensorModule()
        self.final_norm = TensorModule()
        self.lm_head = TensorModule()
        self.layers = {} if load_weights else []
        self.num_layers = num_layers
        if not load_weights:
            return

        self._quant_attrs = quant_attrs
        self._load_quant_config(quant_attrs)

        lm_head_tensors = {}
        for weight_file in os.listdir(input_path):
            if weight_file.endswith(".safetensors"):
                weights = load_file(os.path.join(input_path, weight_file))

                # Map weights to modules
                for raw_name, tensor in weights.items():
                    name = normalize_vlm_weight_name(raw_name)
                    if name is None:
                        continue

                    # Per-layer quantization support
                    local_bits = self.get_layer_bits(name)
                    local_group_size = self.get_layer_group_size(name)

                    if name == "model.embed_tokens.weight" or name == "transformer.embedding.word_embeddings.weight":
                        self.embedding.weight = tensor
                    elif name == "model.norm.weight" or name == "transformer.encoder.final_layernorm.weight":
                        self.final_norm.weight = tensor
                    elif name == "model.norm.bias" or name == "transformer.encoder.final_layernorm.bias":
                        self.final_norm.bias = tensor
                    elif name in {
                        "lm_head.weight",
                        "lm_head.bias",
                        "lm_head.qweight",
                        "lm_head.qzeros",
                        "lm_head.weight_zero_point",
                        "lm_head.scales",
                        "lm_head.weight_scale",
                        "lm_head.g_idx",
                        "transformer.output_layer.weight",
                        "transformer.output_layer.bias",
                        "transformer.output_layer.qweight",
                        "transformer.output_layer.qzeros",
                        "transformer.output_layer.weight_zero_point",
                        "transformer.output_layer.scales",
                        "transformer.output_layer.weight_scale",
                        "transformer.output_layer.g_idx",
                    }:
                        lm_head_tensors[name] = (tensor, local_bits, local_group_size)
                    elif name == "transformer.rotary_pos_emb.inv_freq":
                        # transformer.rotary_pos_emb.inv_freq in ChatGLM3.
                        # Skip rotary embedding weights since they can be re-calculated when looping through the model
                        continue
                    else:
                        if name.startswith("transformer.encoder"):
                            # Chatglm3, e.g., transformer.encoder.layers.0.input_layernorm.weight
                            name = name.replace("transformer.encoder", "model")
                        layer_id = int(name.split(".")[2])
                        module = self.layers.setdefault(layer_id, QuantizedDecoderLayer(layer_id))

                        # Map weights and biases of norm, attention, and feed-forward network
                        # Graph order is input_layernorm --> q_proj/k_proj/v_proj --> o_proj --> post_attention_layernorm --> gate_proj/up_proj --> down_proj
                        # If model uses q_norm and k_norm, graph order is input_layernorm --> q_norm/q_proj/k_norm/k_proj/v_proj --> o_proj --> post_attention_layernorm --> gate_proj/up_proj --> down_proj
                        tensor_map = {}
                        if bool(re.match(r"^model.layers\.\d+\.input_layernorm\.weight$", name)):
                            # model.layers.layer_id.input_layernorm.weight
                            tensor_map["input_layernorm.weight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.input_layernorm\.bias$", name)):
                            # model.layers.layer_id.input_layernorm.bias
                            tensor_map["input_layernorm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.rotary_emb\.inv_freq$", name)):
                            # model.layers.layer_id.self_attn.rotary_emb.inv_freq
                            # Skip rotary embedding weights since they can be re-calculated when looping through the model
                            continue
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.q_proj\.q?weight$", name)):
                            # model.layers.layer_id.self_attn.q_proj.weight
                            # model.layers.layer_id.self_attn.q_proj.qweight
                            tensor_map["self_attn.q_proj.qweight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.q_proj\.(scales|weight_scale)$", name)):
                            # model.layers.layer_id.self_attn.q_proj.scales
                            # model.layers.layer_id.self_attn.q_proj.weight_scale
                            tensor_map["self_attn.q_proj.scales"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.q_proj\.(qzeros|weight_zero_point)$", name)):
                            # model.layers.layer_id.self_attn.q_proj.qzeros
                            # model.layers.layer_id.self_attn.q_proj.weight_zero_point
                            tensor_map["self_attn.q_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.q_proj\.g_idx$", name)):
                            # model.layers.layer_id.self_attn.q_proj.g_idx
                            tensor_map["self_attn.q_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.q_proj\.bias$", name)):
                            # model.layers.layer_id.self_attn.q_proj.bias
                            tensor_map["self_attn.q_proj.bias"] = tensor
                        elif bool(re.match(r"^model\.layers\.\d+\.self_attn\.(q_norm|query_layernorm)\.weight$", name)):
                            tensor_map["self_attn.q_norm.weight"] = tensor
                        elif bool(re.match(r"^model\.layers\.\d+\.self_attn\.(q_norm|query_layernorm)\.bias$", name)):
                            tensor_map["self_attn.q_norm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.k_proj\.q?weight$", name)):
                            # model.layers.layer_id.self_attn.k_proj.qweight
                            # model.layers.layer_id.self_attn.k_proj.weight
                            tensor_map["self_attn.k_proj.qweight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.k_proj\.(scales|weight_scale)$", name)):
                            # model.layers.layer_id.self_attn.k_proj.scales
                            # model.layers.layer_id.self_attn.k_proj.weight_scale
                            tensor_map["self_attn.k_proj.scales"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.k_proj\.(qzeros|weight_zero_point)$", name)):
                            # model.layers.layer_id.self_attn.k_proj.qzeros
                            # model.layers.layer_id.self_attn.k_proj.weight_zero_point
                            tensor_map["self_attn.k_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.k_proj\.g_idx$", name)):
                            # model.layers.layer_id.self_attn.k_proj.g_idx
                            tensor_map["self_attn.k_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.k_proj\.bias$", name)):
                            # model.layers.layer_id.self_attn.k_proj.bias
                            tensor_map["self_attn.k_proj.bias"] = tensor
                        elif bool(re.match(r"^model\.layers\.\d+\.self_attn\.(k_norm|key_layernorm)\.weight$", name)):
                            tensor_map["self_attn.k_norm.weight"] = tensor
                        elif bool(re.match(r"^model\.layers\.\d+\.self_attn\.(k_norm|key_layernorm)\.bias$", name)):
                            tensor_map["self_attn.k_norm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.v_proj\.q?weight$", name)):
                            # model.layers.layer_id.self_attn.v_proj.qweight
                            # model.layers.layer_id.self_attn.v_proj.weight
                            tensor_map["self_attn.v_proj.qweight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.v_proj\.(scales|weight_scale)$", name)):
                            # model.layers.layer_id.self_attn.v_proj.scales
                            # model.layers.layer_id.self_attn.v_proj.weight_scale
                            tensor_map["self_attn.v_proj.scales"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.v_proj\.(qzeros|weight_zero_point)$", name)):
                            # model.layers.layer_id.self_attn.v_proj.qzeros
                            # model.layers.layer_id.self_attn.v_proj.weight_zero_point
                            tensor_map["self_attn.v_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.v_proj\.g_idx$", name)):
                            # model.layers.layer_id.self_attn.v_proj.g_idx
                            tensor_map["self_attn.v_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn.v_proj\.bias$", name)):
                            # model.layers.layer_id.self_attn.v_proj.bias
                            tensor_map["self_attn.v_proj.bias"] = tensor
                        elif bool(
                            re.match(r"^model.layers\.\d+\.(self_attn.o_proj|self_attention.dense)\.q?weight$", name)
                        ):
                            # model.layers.layer_id.self_attn.o_proj.qweight
                            # model.layers.layer_id.self_attention.dense.qweight
                            tensor_map["self_attn.o_proj.qweight"] = tensor
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.o_proj|self_attention.dense)\.(scales|weight_scale)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.self_attn.o_proj.scales
                            # model.layers.layer_id.self_attention.dense.scales
                            # model.layers.layer_id.self_attn.o_proj.weight_scale
                            # model.layers.layer_id.self_attention.dense.weight_scale
                            tensor_map["self_attn.o_proj.scales"] = tensor
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.o_proj|self_attention.dense)\.(qzeros|weight_zero_point)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.self_attn.o_proj.qzeros
                            # model.layers.layer_id.self_attention.dense.qzeros
                            # model.layers.layer_id.self_attn.o_proj.weight_zero_point
                            # model.layers.layer_id.self_attention.dense.weight_zero_point
                            tensor_map["self_attn.o_proj.qzeros"] = tensor
                        elif bool(
                            re.match(r"^model.layers\.\d+\.(self_attn.o_proj|self_attention.dense)\.g_idx$", name)
                        ):
                            # model.layers.layer_id.self_attn.o_proj.g_idx
                            # model.layers.layer_id.self_attention.dense.g_idx
                            tensor_map["self_attn.o_proj.g_idx"] = tensor
                        elif bool(
                            re.match(r"^model.layers\.\d+\.(self_attn.o_proj|self_attention.dense)\.bias$", name)
                        ):
                            # model.layers.layer_id.self_attn.o_proj.bias
                            # model.layers.layer_id.self_attention.dense.bias
                            tensor_map["self_attn.o_proj.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.post_attention_layernorm\.weight$", name)):
                            # model.layers.layer_id.post_attention_layernorm.weight
                            tensor_map["post_attention_layernorm.weight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.post_attention_layernorm\.bias$", name)):
                            # model.layers.layer_id.post_attention_layernorm.bias
                            tensor_map["post_attention_layernorm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.pre_feedforward_layernorm\.weight$", name)):
                            # model.layers.layer_id.pre_feedforward_layernorm.weight
                            tensor_map["pre_feedforward_layernorm.weight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.pre_feedforward_layernorm\.bias$", name)):
                            # model.layers.layer_id.pre_feedforward_layernorm.bias
                            tensor_map["pre_feedforward_layernorm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.post_feedforward_layernorm\.weight$", name)):
                            # model.layers.layer_id.post_feedforward_layernorm.weight
                            tensor_map["post_feedforward_layernorm.weight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.post_feedforward_layernorm\.bias$", name)):
                            # model.layers.layer_id.post_feedforward_layernorm.bias
                            tensor_map["post_feedforward_layernorm.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.gate_proj\.q?weight$", name)):
                            # model.layers.layer_id.mlp.gate_proj.qweight
                            # model.layers.layer_id.mlp.gate_proj.weight
                            tensor_map["mlp.gate_proj.qweight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.gate_proj\.(scales|weight_scale)$", name)):
                            # model.layers.layer_id.mlp.gate_proj.scales
                            # model.layers.layer_id.mlp.gate_proj.weight_scale
                            tensor_map["mlp.gate_proj.scales"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.gate_proj\.(qzeros|weight_zero_point)$", name)):
                            # model.layers.layer_id.mlp.gate_proj.qzeros
                            # model.layers.layer_id.mlp.gate_proj.weight_zero_point
                            tensor_map["mlp.gate_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.gate_proj\.g_idx$", name)):
                            # model.layers.layer_id.mlp.gate_proj.g_idx
                            tensor_map["mlp.gate_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.gate_proj\.bias$", name)):
                            # model.layers.layer_id.mlp.gate_proj.bias
                            tensor_map["mlp.gate_proj.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.up_proj\.q?weight$", name)):
                            # model.layers.layer_id.mlp.up_proj.qweight
                            # model.layers.layer_id.mlp.up_proj.weight
                            tensor_map["mlp.up_proj.qweight"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.up_proj\.(scales|weight_scale)$", name)):
                            # model.layers.layer_id.mlp.up_proj.scales
                            # model.layers.layer_id.mlp.up_proj.weight_scale
                            tensor_map["mlp.up_proj.scales"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.up_proj\.(qzeros|weight_zero_point)$", name)):
                            # model.layers.layer_id.mlp.up_proj.qzeros
                            # model.layers.layer_id.mlp.up_proj.weight_zero_point
                            tensor_map["mlp.up_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.up_proj\.g_idx$", name)):
                            # model.layers.layer_id.mlp.up_proj.g_idx
                            tensor_map["mlp.up_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.up_proj\.bias$", name)):
                            # model.layers.layer_id.mlp.up_proj.bias
                            tensor_map["mlp.up_proj.bias"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.(down_proj|dense_4h_to_h)\.q?weight$", name)):
                            # model.layers.layer_id.mlp.down_proj.qweight
                            # model.layers.layer_id.mlp.dense_4h_to_h.qweight
                            # model.layers.layer_id.mlp.down_proj.weight
                            # model.layers.layer_id.mlp.dense_4h_to_h.weight
                            tensor_map["mlp.down_proj.qweight"] = tensor
                        elif bool(
                            re.match(r"^model.layers\.\d+\.mlp.(down_proj|dense_4h_to_h)\.(scales|weight_scale)$", name)
                        ):
                            # model.layers.layer_id.mlp.down_proj.scales
                            # model.layers.layer_id.mlp.dense_4h_to_h.scales
                            # model.layers.layer_id.mlp.down_proj.weight_scale
                            # model.layers.layer_id.mlp.dense_4h_to_h.weight_scale
                            tensor_map["mlp.down_proj.scales"] = tensor
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.mlp.(down_proj|dense_4h_to_h)\.(qzeros|weight_zero_point)$", name
                            )
                        ):
                            # model.layers.layer_id.mlp.down_proj.qzeros
                            # model.layers.layer_id.mlp.dense_4h_to_h.qzeros
                            # model.layers.layer_id.mlp.down_proj.weight_zero_point
                            # model.layers.layer_id.mlp.dense_4h_to_h.weight_zero_point
                            tensor_map["mlp.down_proj.qzeros"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.(down_proj|dense_4h_to_h)\.g_idx$", name)):
                            # model.layers.layer_id.mlp.down_proj.g_idx
                            # model.layers.layer_id.mlp.dense_4h_to_h.g_idx
                            tensor_map["mlp.down_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.(down_proj|dense_4h_to_h)\.bias$", name)):
                            # model.layers.layer_id.mlp.down_proj.bias
                            # model.layers.layer_id.mlp.dense_4h_to_h.bias
                            tensor_map["mlp.down_proj.bias"] = tensor
                        # Match against fused layers
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.q?weight$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.self_attn.qkv_proj.qweight
                            # model.layers.layer_id.self_attention.query_key_value.qweight
                            # model.layers.layer_id.self_attn.qkv_proj.weight
                            # model.layers.layer_id.self_attention.query_key_value.weight
                            if quant_type == "olive":
                                # Olive: (out_features, in_features), split on dim=0
                                tensor_map["self_attn.q_proj.qweight"] = tensor[:q_size, :]
                                tensor_map["self_attn.k_proj.qweight"] = tensor[q_size : q_size + kv_size, :]
                                tensor_map["self_attn.v_proj.qweight"] = tensor[q_size + kv_size :, :]
                            else:
                                # AWQ/GPTQ/Quark: (in_features, out_features), split on dim=1
                                q_dim = q_size // (32 // local_bits) if quant_type in {"awq", "quark"} else q_size
                                kv_dim = kv_size // (32 // local_bits) if quant_type in {"awq", "quark"} else kv_size
                                tensor_map["self_attn.q_proj.qweight"] = tensor[:, :q_dim]
                                tensor_map["self_attn.k_proj.qweight"] = tensor[:, q_dim : q_dim + kv_dim]
                                tensor_map["self_attn.v_proj.qweight"] = tensor[:, q_dim + kv_dim :]
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.(scales|weight_scale)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.self_attn.qkv_proj.scales
                            # model.layers.layer_id.self_attention.query_key_value.scales
                            # model.layers.layer_id.self_attn.qkv_proj.weight_scale
                            # model.layers.layer_id.self_attention.query_key_value.weight_scale
                            if quant_type == "olive":
                                # Olive: (out_features, num_groups), split on dim=0
                                tensor_map["self_attn.q_proj.scales"] = tensor[:q_size, :]
                                tensor_map["self_attn.k_proj.scales"] = tensor[q_size : q_size + kv_size, :]
                                tensor_map["self_attn.v_proj.scales"] = tensor[q_size + kv_size :, :]
                            else:
                                # AWQ/GPTQ/Quark: split on dim=1
                                tensor_map["self_attn.q_proj.scales"] = tensor[:, :q_size]
                                tensor_map["self_attn.k_proj.scales"] = tensor[:, q_size : q_size + kv_size]
                                tensor_map["self_attn.v_proj.scales"] = tensor[:, q_size + kv_size :]
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.(qzeros|weight_zero_point)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.self_attn.qkv_proj.qzeros
                            # model.layers.layer_id.self_attention.query_key_value.qzeros
                            # model.layers.layer_id.self_attn.qkv_proj.weight_zero_point
                            # model.layers.layer_id.self_attention.query_key_value.weight_zero_point
                            if quant_type == "olive":
                                # Olive: (out_features, packed_num_groups) uint8, split on dim=0
                                q_dim = q_size // (8 // local_bits)
                                kv_dim = kv_size // (8 // local_bits)
                                tensor_map["self_attn.q_proj.qzeros"] = tensor[:q_dim, :]
                                tensor_map["self_attn.k_proj.qzeros"] = tensor[q_dim : q_dim + kv_dim, :]
                                tensor_map["self_attn.v_proj.qzeros"] = tensor[q_dim + kv_dim :, :]
                            else:
                                # AWQ/GPTQ/Quark: int32 packing, split on dim=1
                                q_dim = (
                                    q_size // (32 // local_bits)
                                    if quant_type in {"awq", "gptq", "quark"}
                                    else q_size
                                )
                                kv_dim = (
                                    kv_size // (32 // local_bits)
                                    if quant_type in {"awq", "gptq", "quark"}
                                    else kv_size
                                )
                                tensor_map["self_attn.q_proj.qzeros"] = tensor[:, :q_dim]
                                tensor_map["self_attn.k_proj.qzeros"] = tensor[:, q_dim : q_dim + kv_dim]
                                tensor_map["self_attn.v_proj.qzeros"] = tensor[:, q_dim + kv_dim :]
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.g_idx$", name
                            )
                        ):
                            # model.layers.layer_id.self_attn.qkv_proj.g_idx
                            # model.layers.layer_id.self_attention.query_key_value.g_idx
                            tensor_map["self_attn.q_proj.g_idx"] = tensor
                            tensor_map["self_attn.k_proj.g_idx"] = tensor
                            tensor_map["self_attn.v_proj.g_idx"] = tensor
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.(self_attn.qkv_proj|self_attention.query_key_value)\.bias$", name
                            )
                        ):
                            # model.layers.layer_id.self_attn.qkv_proj.bias
                            # model.layers.layer_id.self_attention.query_key_value.bias
                            tensor_map["self_attn.q_proj.bias"] = tensor[:q_size]
                            tensor_map["self_attn.k_proj.bias"] = tensor[q_size : q_size + kv_size]
                            tensor_map["self_attn.v_proj.bias"] = tensor[q_size + kv_size :]
                        elif bool(
                            re.match(r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h|gate_proj)\.q?weight$", name)
                        ):
                            # model.layers.layer_id.mlp.gate_up_proj.qweight
                            # model.layers.layer_id.mlp.dense_h_to_4h.qweight
                            # model.layers.layer_id.mlp.gate_up_proj.weight
                            # model.layers.layer_id.mlp.dense_h_to_4h.weight
                            if quant_type == "olive":
                                # Olive: (out_features, in_features), split on dim=0
                                tensor_map["mlp.gate_proj.qweight"] = tensor[:intermediate_size, :]
                                tensor_map["mlp.up_proj.qweight"] = tensor[intermediate_size:, :]
                            else:
                                # AWQ/GPTQ/Quark: (in_features, out_features), split on dim=1
                                intermediate_dim = (
                                    intermediate_size // (32 // local_bits)
                                    if quant_type in {"awq", "quark"}
                                    else intermediate_size
                                )
                                tensor_map["mlp.gate_proj.qweight"] = tensor[:, :intermediate_dim]
                                tensor_map["mlp.up_proj.qweight"] = tensor[:, intermediate_dim:]
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h|gate_proj)\.(scales|weight_scale)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.mlp.gate_up_proj.scales
                            # model.layers.layer_id.mlp.dense_h_to_4h.scales
                            # model.layers.layer_id.mlp.gate_up_proj.weight_scale
                            # model.layers.layer_id.mlp.dense_h_to_4h.weight_scale
                            if quant_type == "olive":
                                # Olive: (out_features, num_groups), split on dim=0
                                tensor_map["mlp.gate_proj.scales"] = tensor[:intermediate_size, :]
                                tensor_map["mlp.up_proj.scales"] = tensor[intermediate_size:, :]
                            else:
                                # AWQ/GPTQ/Quark: split on dim=1
                                tensor_map["mlp.gate_proj.scales"] = tensor[:, :intermediate_size]
                                tensor_map["mlp.up_proj.scales"] = tensor[:, intermediate_size:]
                        elif bool(
                            re.match(
                                r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h|gate_proj)\.(qzeros|weight_zero_point)$",
                                name,
                            )
                        ):
                            # model.layers.layer_id.mlp.gate_up_proj.qzeros
                            # model.layers.layer_id.mlp.dense_h_to_4h.qzeros
                            # model.layers.layer_id.mlp.gate_up_proj.weight_zero_point
                            # model.layers.layer_id.mlp.dense_h_to_4h.weight_zero_point
                            if quant_type == "olive":
                                # Olive: (out_features, packed_num_groups) uint8, split on dim=0
                                intermediate_dim = intermediate_size // (8 // local_bits)
                                tensor_map["mlp.gate_proj.qzeros"] = tensor[:intermediate_dim, :]
                                tensor_map["mlp.up_proj.qzeros"] = tensor[intermediate_dim:, :]
                            else:
                                # AWQ/GPTQ/Quark: int32 packing, split on dim=1
                                intermediate_dim = (
                                    intermediate_size // (32 // local_bits)
                                    if quant_type in {"awq", "gptq", "quark"}
                                    else intermediate_size
                                )
                                tensor_map["mlp.gate_proj.qzeros"] = tensor[:, :intermediate_dim]
                                tensor_map["mlp.up_proj.qzeros"] = tensor[:, intermediate_dim:]
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h)\.g_idx$", name)):
                            # model.layers.layer_id.mlp.gate_up_proj.g_idx
                            # model.layers.layer_id.mlp.dense_h_to_4h.g_idx
                            tensor_map["mlp.gate_proj.g_idx"] = tensor
                            tensor_map["mlp.up_proj.g_idx"] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.mlp.(gate_up_proj|dense_h_to_4h)\.bias$", name)):
                            # model.layers.layer_id.mlp.gate_up_proj.bias
                            # model.layers.layer_id.mlp.dense_h_to_4h.bias
                            tensor_map["mlp.gate_proj.bias"] = tensor[:intermediate_size]
                            tensor_map["mlp.up_proj.bias"] = tensor[intermediate_size:]
                        elif bool(re.match(r"^model\.layers\.\d+\.mlp\.experts\.\d+\.(gate_proj|up_proj|gate_up_proj|down_proj)\.(weight|bias|qweight|scales|qzeros|weight_scale|weight_zero_point|g_idx)$", name)):
                            # model.layers.layer_id.mlp.experts.expert_id.proj_type.param_type
                            split_name = name.split(".")
                            expert_id = int(split_name[5])
                            proj_type = split_name[-2]
                            param_type = split_name[-1]
                            module.mlp.experts.set_weight_data(expert_id, proj_type, param_type, tensor, local_bits, local_group_size)
                        elif bool(re.match(r"^model\.layers\.\d+\.mlp\.router\.(weight|bias)$", name)):
                            # model.layers.layer_id.mlp.router.weight
                            # model.layers.layer_id.mlp.router.bias
                            tensor_map["mlp.router." + name.split(".")[-1]] = tensor
                        elif bool(re.match(r"^model.layers\.\d+\.self_attn\.sinks$", name)):
                            # model.layers.layer_id.self_attn.sinks
                            tensor_map["self_attn.sinks"] = tensor
                        else:
                            raise NotImplementedError(f"{name} in your quantized model is not recognized.")

                        for tensor_name, tensor_value in tensor_map.items():
                            submodule = module
                            for sub_name in tensor_name.split(".")[:-1]:
                                submodule = getattr(submodule, sub_name)
                            if isinstance(submodule, QuantizedTensorModule):
                                for q_attr, q_value in [("bits", local_bits), ("_group_size", local_group_size)]:
                                    if getattr(submodule, q_attr) is not None and getattr(submodule, q_attr) != q_value:
                                        raise ValueError(
                                            f"Quantization {q_attr} mismatch for {name}: expected {getattr(submodule, q_attr)}, got {q_value}."
                                        )
                                    setattr(submodule, q_attr, q_value)
                            setattr(submodule, tensor_name.split(".")[-1], tensor_value)

        # Process collected lm_head tensors in defined order to avoid ordering issues
        self._assign_lm_head_tensors(lm_head_tensors)

        # Set LM head weights + biases if not already set
        if isinstance(self.lm_head, TensorModule) and self.lm_head.weight is None:
            # Embedding and LM head share same weights + biases (lm_head.weight == embedding.weight and lm_head.bias == embedding.bias)
            self.lm_head.weight = self.embedding.weight
            if self.lm_head.bias is not None:
                self.lm_head.bias = self.embedding.bias

        # Sort list of layers by layer id
        self.layers = list(self.layers.values())
        self.layers.sort(key=lambda m: m.layer_id)

        # Set properties of each layer based on quantization type
        self.set_properties()

    # Canonical name mapping for lm_head tensors (transformer.output_layer.* -> lm_head.*)
    _LM_HEAD_NAME_MAP = {
        "transformer.output_layer.weight": "lm_head.weight",
        "transformer.output_layer.bias": "lm_head.bias",
        "transformer.output_layer.qweight": "lm_head.qweight",
        "transformer.output_layer.qzeros": "lm_head.qzeros",
        "transformer.output_layer.weight_zero_point": "lm_head.weight_zero_point",
        "transformer.output_layer.scales": "lm_head.scales",
        "transformer.output_layer.weight_scale": "lm_head.weight_scale",
        "transformer.output_layer.g_idx": "lm_head.g_idx",
    }

    def _assign_lm_head_tensors(self, lm_head_tensors):
        """Assign collected lm_head tensors in a defined order so that weight/bias
        are always processed before quantization parameters (scales, qzeros, etc.),
        regardless of safetensors dict iteration order."""
        # Normalize names to canonical lm_head.* form
        normalized = {}
        for name, value in lm_head_tensors.items():
            canonical = self._LM_HEAD_NAME_MAP.get(name, name)
            normalized[canonical] = value

        # Process weight and bias first, then quantization tensors
        ordered_keys = [
            "lm_head.weight",
            "lm_head.bias",
            "lm_head.qweight",
            "lm_head.qzeros",
            "lm_head.weight_zero_point",
            "lm_head.scales",
            "lm_head.weight_scale",
            "lm_head.g_idx",
        ]

        for key in ordered_keys:
            if key not in normalized:
                continue
            tensor, local_bits, local_group_size = normalized[key]
            if key == "lm_head.weight":
                self.lm_head.weight = tensor
            elif key == "lm_head.bias":
                self.lm_head.bias = tensor
            elif key == "lm_head.qweight":
                self._initialize_quantized_lm_head(local_bits, local_group_size)
                self.lm_head.qweight = tensor
            elif key in {"lm_head.qzeros", "lm_head.weight_zero_point"}:
                self._initialize_quantized_lm_head(local_bits, local_group_size)
                self.lm_head.qzeros = tensor
            elif key in {"lm_head.scales", "lm_head.weight_scale"}:
                self._initialize_quantized_lm_head(local_bits, local_group_size)
                self.lm_head.scales = tensor
            elif key == "lm_head.g_idx":
                self._initialize_quantized_lm_head(local_bits, local_group_size)
                self.lm_head.g_idx = tensor

    def _load_quant_config(self, quant_attrs):
        self.global_group_size = quant_attrs["config"]["group_size"]
        self.global_bits = quant_attrs["config"]["bits"]

    def get_layer_bits(self, layer_name):
        # 'bits' is globally defined for all layers
        return self.global_bits

    def get_layer_group_size(self, layer_name):
        # 'group_size' is globally defined for all layers
        return self.global_group_size

    def _initialize_quantized_lm_head(self, bits, group_size):
        """
        Initialize `QuantizedTensorModule` for LM head if not already set
        """
        if not isinstance(self.lm_head, QuantizedTensorModule):
            q_lm_head = QuantizedTensorModule()
            q_lm_head.qweight = self.lm_head.weight
            q_lm_head.bias = self.lm_head.bias
            q_lm_head.bits = bits
            q_lm_head.group_size = group_size
            self.lm_head = q_lm_head

    def set_g_idx(self, module):
        if module is not None and module.g_idx is None:
            module.g_idx = torch.tensor([i // module.group_size for i in range(module.in_features)], dtype=torch.int32)

    def set_properties(self):
        """
        Set in_features, out_features, and g_idx based on quantization type
        """
        if isinstance(self.lm_head, QuantizedTensorModule):
            if self.quant_type == "awq" or self.quant_type == "quark":
                self.lm_head.out_features = self.lm_head.scales.shape[1]
                self.lm_head.in_features = self.lm_head.qweight.shape[0]
                # Set g_idx if not already set
                self.set_g_idx(self.lm_head.g_idx)
            elif self.quant_type == "gptq":
                self.lm_head.out_features = self.lm_head.qweight.shape[1]
                self.lm_head.in_features = self.lm_head.g_idx.shape[0]
            elif self.quant_type == "olive":
                # Olive format: qweight is (out_features, packed_in_features) uint8
                # packed_in_features = in_features * bits / 8
                self.lm_head.out_features = self.lm_head.qweight.shape[0]
                self.lm_head.in_features = self.lm_head.qweight.shape[1] * 8 // self.lm_head.bits
            else:
                raise NotImplementedError(f"The {self.quant_type} quantization method is not recognized.")
        for module in self.layers:
            if self.quant_type == "awq" or self.quant_type == "quark":
                # Set in_features and out_features
                module.self_attn.q_proj.out_features = module.self_attn.q_proj.scales.shape[1]
                module.self_attn.q_proj.in_features = module.self_attn.q_proj.qweight.shape[0]
                module.self_attn.k_proj.out_features = module.self_attn.k_proj.scales.shape[1]
                module.self_attn.k_proj.in_features = module.self_attn.k_proj.qweight.shape[0]
                module.self_attn.v_proj.out_features = module.self_attn.v_proj.scales.shape[1]
                module.self_attn.v_proj.in_features = module.self_attn.v_proj.qweight.shape[0]
                module.self_attn.o_proj.out_features = module.self_attn.o_proj.scales.shape[1]
                module.self_attn.o_proj.in_features = module.self_attn.o_proj.qweight.shape[0]

                # Set g_idx if not already set
                self.set_g_idx(module.self_attn.q_proj.g_idx)
                self.set_g_idx(module.self_attn.k_proj.g_idx)
                self.set_g_idx(module.self_attn.v_proj.g_idx)
                self.set_g_idx(module.self_attn.o_proj.g_idx)
                self.set_g_idx(module.mlp.gate_proj.g_idx)
                self.set_g_idx(module.mlp.up_proj.g_idx)
                self.set_g_idx(module.mlp.down_proj.g_idx)

                if module.mlp.experts.num_experts > 0:
                    for _, expert in module.mlp.experts.items():
                        if expert.gate_up_proj.qweight is not None:
                            expert.gate_up_proj.out_features = expert.gate_up_proj.scales.shape[1]
                            expert.gate_up_proj.in_features = expert.gate_up_proj.qweight.shape[0]
                            # Set g_idx if not already set
                            self.set_g_idx(expert.gate_up_proj.g_idx)
                            self.set_g_idx(expert.down_proj.g_idx)
                        else:
                            expert.gate_proj.out_features = expert.gate_proj.scales.shape[1]
                            expert.gate_proj.in_features = expert.gate_proj.qweight.shape[0]
                            expert.up_proj.out_features = expert.up_proj.scales.shape[1]
                            expert.up_proj.in_features = expert.up_proj.qweight.shape[0]
                            # Set g_idx if not already set
                            self.set_g_idx(expert.gate_proj.g_idx)
                            self.set_g_idx(expert.up_proj.g_idx)
                            self.set_g_idx(expert.down_proj.g_idx)

                        expert.down_proj.out_features = expert.down_proj.scales.shape[1]
                        expert.down_proj.in_features = expert.down_proj.qweight.shape[0]
                else:
                    module.mlp.gate_proj.out_features = module.mlp.gate_proj.scales.shape[1]
                    module.mlp.gate_proj.in_features = module.mlp.gate_proj.qweight.shape[0]
                    module.mlp.up_proj.out_features = module.mlp.up_proj.scales.shape[1]
                    module.mlp.up_proj.in_features = module.mlp.up_proj.qweight.shape[0]
                    module.mlp.down_proj.out_features = module.mlp.down_proj.scales.shape[1]
                    module.mlp.down_proj.in_features = module.mlp.down_proj.qweight.shape[0]

                    self.set_g_idx(module.mlp.gate_proj.g_idx)
                    self.set_g_idx(module.mlp.up_proj.g_idx)
                    self.set_g_idx(module.mlp.down_proj.g_idx)
            elif self.quant_type == "gptq":
                # Set in_features and out_features
                module.self_attn.q_proj.out_features = module.self_attn.q_proj.qweight.shape[1]
                module.self_attn.q_proj.in_features = module.self_attn.q_proj.g_idx.shape[0]
                module.self_attn.k_proj.out_features = module.self_attn.k_proj.qweight.shape[1]
                module.self_attn.k_proj.in_features = module.self_attn.k_proj.g_idx.shape[0]
                module.self_attn.v_proj.out_features = module.self_attn.v_proj.qweight.shape[1]
                module.self_attn.v_proj.in_features = module.self_attn.v_proj.g_idx.shape[0]
                module.self_attn.o_proj.out_features = module.self_attn.o_proj.qweight.shape[1]
                module.self_attn.o_proj.in_features = module.self_attn.o_proj.g_idx.shape[0]
                module.mlp.gate_proj.out_features = module.mlp.gate_proj.qweight.shape[1]
                module.mlp.gate_proj.in_features = module.mlp.gate_proj.g_idx.shape[0]
                module.mlp.up_proj.out_features = module.mlp.up_proj.qweight.shape[1]
                module.mlp.up_proj.in_features = module.mlp.up_proj.g_idx.shape[0]
                module.mlp.down_proj.out_features = module.mlp.down_proj.qweight.shape[1]
                module.mlp.down_proj.in_features = module.mlp.down_proj.g_idx.shape[0]

            elif self.quant_type == "olive":
                module.self_attn.q_proj.out_features = module.self_attn.q_proj.qweight.shape[0]
                module.self_attn.q_proj.in_features = (
                    module.self_attn.q_proj.qweight.shape[1] * 8 // module.self_attn.q_proj.bits
                )
                module.self_attn.k_proj.out_features = module.self_attn.k_proj.qweight.shape[0]
                module.self_attn.k_proj.in_features = (
                    module.self_attn.k_proj.qweight.shape[1] * 8 // module.self_attn.k_proj.bits
                )
                module.self_attn.v_proj.out_features = module.self_attn.v_proj.qweight.shape[0]
                module.self_attn.v_proj.in_features = (
                    module.self_attn.v_proj.qweight.shape[1] * 8 // module.self_attn.v_proj.bits
                )
                module.self_attn.o_proj.out_features = module.self_attn.o_proj.qweight.shape[0]
                module.self_attn.o_proj.in_features = (
                    module.self_attn.o_proj.qweight.shape[1] * 8 // module.self_attn.o_proj.bits
                )
                module.mlp.gate_proj.out_features = module.mlp.gate_proj.qweight.shape[0]
                module.mlp.gate_proj.in_features = (
                    module.mlp.gate_proj.qweight.shape[1] * 8 // module.mlp.gate_proj.bits
                )
                module.mlp.up_proj.out_features = module.mlp.up_proj.qweight.shape[0]
                module.mlp.up_proj.in_features = module.mlp.up_proj.qweight.shape[1] * 8 // module.mlp.up_proj.bits
                module.mlp.down_proj.out_features = module.mlp.down_proj.qweight.shape[0]
                module.mlp.down_proj.in_features = (
                    module.mlp.down_proj.qweight.shape[1] * 8 // module.mlp.down_proj.bits
                )

            else:
                raise NotImplementedError(f"The {self.quant_type} quantization method is not recognized.")

    def modules(self):
        """
        Return list of modules in quantized model in order of appearance in the model
        """
        return [self.embedding] + self.layers + [self.final_norm, self.lm_head]

    def unpack(self, module):
        """
        Unpack `qzeros` and `qweight` to standard format
        """
        self.unpack_qzeros(module)
        self.unpack_qweight(module)
        self.dequant_weight(module)

    def repack(self, module):
        """
        Repack `scales`, `qzeros` and `qweight` to ORT format
        """
        intweight = self.quant_weight(module)
        self.pack_ort_format(module, intweight)

    def unpack_qzeros(self, module):
        """
        Unpack `qzeros` to standard format
        """
        if module.qzeros is None:
            return
        expected_shape = (module.in_features // module.group_size, module.out_features)
        transpose = module.qzeros.shape[0] != expected_shape[0]
        module.qzeros = self.unpack_on_row(module.qzeros, module.bits, transpose)

    def unpack_qweight(self, module):
        """
        Unpack `qweight` to standard format
        """
        expected_shape = (module.in_features, module.qweight.shape[1])
        transpose = module.qweight.shape[0] != expected_shape[0]
        module.qweight = self.unpack_on_row(module.qweight, module.bits, transpose)

    def pack_qzeros(self, module):
        """
        Pack `qzeros` to quantized format
        """
        expected_shape = (module.in_features // module.group_size, module.out_features)
        transpose = module.qzeros.shape[0] != expected_shape[0]
        module.qzeros = self.pack_on_row(module.qzeros, module.bits, transpose)

    def unpack_on_row_for_2_4_8_bits(self, tensor, bits, transpose):
        """
        Perform general-purpose unpacking on 2-bit, 4-bit, or 8-bit tensor
        """
        pack_tensor = tensor.T if transpose else tensor
        wf = torch.arange(0, 32, bits, device=pack_tensor.device).unsqueeze(0).unsqueeze(0)
        out = torch.bitwise_right_shift(torch.unsqueeze(pack_tensor, 2), wf)
        out = out.reshape(pack_tensor.shape[0], -1)
        out = torch.bitwise_and(out, (2**bits) - 1)
        return out.T if transpose else out

    def unpack_on_row(self, tensor, bits, transpose):
        """
        Unpack tensor by row. Packed datatype is assumed to be int32.
        """
        if bits in {2, 4, 8}:
            return self.unpack_on_row_for_2_4_8_bits(tensor, bits, transpose)
        else:
            raise NotImplementedError(f"Unpacking for {bits}-bit quantization is not currently supported.")

    def pack_on_row_for_2_4_8_bits(self, tensor, bits, transpose, packed_dtype=torch.int32):
        """
        Perform general-purpose packing on 2-bit, 4-bit, or 8-bit tensor
        """
        packed_bitwidth = torch.iinfo(packed_dtype).bits
        values_per_pack = packed_bitwidth // bits

        orig_tensor = tensor.T if transpose else tensor

        original_cols = orig_tensor.shape[1]
        pad_len = (values_per_pack - (original_cols % values_per_pack)) % values_per_pack
        if pad_len > 0:
            orig_tensor = torch.nn.functional.pad(orig_tensor, (0, pad_len), "constant", 0)

        wf = torch.arange(0, bits).view(1, 1, -1)
        out = torch.bitwise_right_shift(orig_tensor.unsqueeze(-1), wf)
        out = torch.bitwise_and(out, 1)

        out = out.reshape(orig_tensor.shape[0], -1, values_per_pack * bits)
        wf1 = torch.arange(0, values_per_pack * bits, 1).view(1, 1, -1)
        out = torch.bitwise_left_shift(out, wf1)
        out = out.sum(dim=-1).to(packed_dtype)
        return out.T if transpose else out

    def pack_on_row(self, tensor, bits, transpose, packed_dtype=torch.int32):
        """
        Pack tensor by row
        """
        if bits in {2, 4, 8}:
            return self.pack_on_row_for_2_4_8_bits(tensor, bits, transpose, packed_dtype)
        else:
            raise NotImplementedError(f"Packing for {bits}-bit quantization is not currently supported.")

    def dequant_weight(self, module):
        """
        De-quantize `qweight` to higher precision (float16)
        """
        # Note: `qweight` and `qzeros` have already been unpacked and stored in those variables respectively
        intweight = module.qweight
        zeros = module.qzeros
        scales = module.scales
        g_idx = module.g_idx

        # De-quantize weight to higher precision
        scale_zeros = zeros * scales
        if g_idx is not None:
            scales = scales[g_idx]
            scale_zeros = scale_zeros[g_idx]
        elif module.group_size != module.in_features:
            scales = scales.repeat_interleave(module.group_size, 0)
            scale_zeros = scale_zeros.repeat_interleave(module.group_size, 0)
        qdq_weight_T = intweight * scales - scale_zeros.half()

        # Store unpacked result in `qweight`
        module.qweight = qdq_weight_T.T

    def quant_weight(self, module):
        """
        Calculate integer weight to quantize `qweight` with
        """
        weight = module.qweight.T
        zeros = module.qzeros
        scales = module.scales
        g_idx = module.g_idx

        scale_zeros = zeros * scales
        if g_idx is not None:
            scales = scales[g_idx]
            scale_zeros = scale_zeros[g_idx]
        elif module.group_size != module.in_features:
            scales = scales.repeat_interleave(module.group_size, 0)
            scale_zeros = scale_zeros.repeat_interleave(module.group_size, 0)
        intweight_T = torch.round((weight + scale_zeros) / scales).to(torch.int)

        return intweight_T

    def pack_ort_format(self, module, intweight):
        """
        Pack `scales`, `qzeros`, and `qweight` to ORT format
        """
        if module.bits not in [2, 4, 8]:
            raise NotImplementedError(f"{module.bits}-bit quantization in ORT is not currently supported by this tool.")

        intweight_pt = intweight.byte()
        kpack = 8 // module.bits
        block_size = module.group_size

        rows, cols = intweight_pt.shape
        blob_size = (block_size + kpack - 1) // kpack
        k_blocks = (rows + block_size - 1) // block_size
        padded_rows = k_blocks * block_size
        pad_len = padded_rows - rows
        if pad_len > 0:
            intweight_pt = torch.nn.functional.pad(intweight_pt, (0, 0, 0, pad_len), "constant", 0)

        intweight_pt_T = intweight.T
        intweight_pt_T = self.pack_on_row(intweight_pt_T, module.bits, transpose=False, packed_dtype=torch.uint8)
        intweight_pt_T = intweight_pt_T.reshape(cols, k_blocks, blob_size)

        scales_pt = module.scales.T.reshape(-1)

        module.scales = scales_pt.contiguous()
        module.qweight = intweight_pt_T.contiguous().byte()

        self.pack_zeros_ort_format(module, reshape=True)

    def pack_zeros_ort_format(self, module, reshape=False):
        """
        Pack `qzeros` to ORT format
        """
        if module.bits not in [2, 4, 8]:
            raise NotImplementedError(f"{module.bits}-bit quantization in ORT is not currently supported by this tool.")

        intzeros_pt = module.qzeros.T if module.qzeros.dtype == module.scales.dtype else module.qzeros.T.byte()

        if module.qzeros.dtype != module.scales.dtype:
            intzeros_pt = self.pack_on_row(intzeros_pt, module.bits, transpose=False, packed_dtype=torch.uint8)
            if reshape:
                intzeros_pt = intzeros_pt.reshape(-1)

        if module.qzeros.dtype != module.scales.dtype:
            module.qzeros = intzeros_pt.contiguous().byte()
        else:
            module.qzeros = intzeros_pt.contiguous()


class AWQModel(QuantizedModel):
    def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers):
        super().__init__(quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers)

        # Unpack and repack all `QuantizedTensorModule` classes in model
        for i, layer in enumerate(self.layers):
            if i >= self.num_layers:
                break
            print(f"Unpacking and repacking layer {i}")

            # Unpack and repack all `QuantizedTensorModule` classes in attention
            self_attn = getattr(layer, "self_attn", None) or getattr(layer, "self_attention", None)
            for _, q_tensors in self_attn.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    # Set `g_idx` to None since it's not used in `MatMulNBits`
                    q_tensors.g_idx = None

            # Unpack and repack all `QuantizedTensorModule` classes in MLP
            for _, q_tensors in layer.mlp.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    # Set `g_idx` to None since it's not used in `MatMulNBits`
                    q_tensors.g_idx = None

        if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None:
            self.unpack(self.lm_head)
            self.repack(self.lm_head)

            # Set `g_idx` to None since it's not used in `MatMulNBits`
            self.lm_head.g_idx = None

    def unpack_qweight(self, module):
        """
        Unpack `qweight` to standard format
        """
        expected_shape = (module.qweight.shape[0], module.out_features)
        transpose = module.qweight.shape != expected_shape
        module.qweight = self.unpack_on_row(module.qweight.T, module.bits, transpose)
        module.qweight = self.reverse_reorder_tensor(module.qweight.T, module.bits)

    def unpack_qzeros(self, module):
        """
        Unpack `qzeros` to standard format
        """
        super().unpack_qzeros(module)
        module.qzeros = self.reverse_reorder_tensor(module.qzeros, module.bits)

    def reverse_reorder_tensor(self, tensor, bits):
        """
        Re-arrange tensor data in a new order
        """
        compress_ratio = 32 // bits
        assert tensor.shape[-1] % compress_ratio == 0

        if bits == 4:
            order_map = [0, 2, 4, 6, 1, 3, 5, 7]
        else:
            raise NotImplementedError(f"Unpacking for {bits}-bit quantization is not currently supported.")

        order_tensor = torch.tensor(order_map, dtype=torch.int32).reshape(1, -1)
        order_tensor = order_tensor.repeat(tensor.shape[1] // compress_ratio, 1)
        order_tensor = order_tensor + torch.arange(0, tensor.shape[1], compress_ratio, dtype=torch.int32).reshape(-1, 1)
        order_tensor = order_tensor.reshape(-1)

        reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
        reverse_order_tensor = reverse_order_tensor[order_tensor]
        int_tensor = tensor[:, reverse_order_tensor]
        return int_tensor


class GPTQModel(QuantizedModel):
    def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers):
        super().__init__(quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers)

        # Unpack and repack all `QuantizedTensorModule` classes in model
        for i, layer in enumerate(self.layers):
            if i >= self.num_layers:
                break
            print(f"Unpacking and repacking layer {i}")

            # Unpack and repack all `QuantizedTensorModule` classes in attention
            for _, q_tensors in layer.self_attn.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.handle_qzeros(q_tensors)
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    if not quant_attrs["use_g_idx"]:
                        # Set `g_idx` to None since it's not used in `MatMulNBits`
                        q_tensors.g_idx = None

            # Unpack and repack all `QuantizedTensorModule` classes in MLP
            for _, q_tensors in layer.mlp.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.handle_qzeros(q_tensors)
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    if not quant_attrs["use_g_idx"]:
                        # Set `g_idx` to None since it's not used in `MatMulNBits`
                        q_tensors.g_idx = None

        if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None:
            self.handle_qzeros(self.lm_head)
            self.unpack(self.lm_head)
            self.repack(self.lm_head)

            if not quant_attrs["use_g_idx"]:
                # Set `g_idx` to None since it's not used in `MatMulNBits`
                self.lm_head.g_idx = None

    def handle_qzeros(self, module):
        """
        Re-pack `qzeros` to handle extra `-1`s
        """
        if module.qzeros is None or module.qzeros.numel() == 0:
            return

        class TempModule:
            def __init__(self, module):
                self.in_features = module.in_features
                self.out_features = module.out_features
                self.group_size = module.group_size
                self.bits = module.bits
                self.qzeros = module.qzeros

        temp_module = TempModule(module)
        self.unpack_qzeros(temp_module)

        temp_module.qzeros += 1
        temp_module.qzeros = torch.bitwise_and(temp_module.qzeros, (2**temp_module.bits) - 1)

        self.pack_qzeros(temp_module)
        module.qzeros = temp_module.qzeros

    def _load_quant_config(self, quant_attrs):
        super()._load_quant_config(quant_attrs)
        self.overrides = quant_attrs["config"].get("dynamic", {})

    def get_overrides(self, layer_name):
        for pattern, overrides in self.overrides.items():
            if re.match(pattern.removeprefix("+:"), layer_name):
                return overrides
        return {}

    def get_layer_bits(self, layer_name):
        return self.get_overrides(layer_name).get("bits", self.global_bits)

    def get_layer_group_size(self, layer_name):
        return self.get_overrides(layer_name).get("group_size", self.global_group_size)


class QuarkModel(QuantizedModel):
    def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers):
        super().__init__(quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers)

        # Unpack and repack all `QuantizedTensorModule` classes in model
        for i, layer in enumerate(self.layers):
            if i >= self.num_layers:
                break
            print(f"Unpacking and repacking layer {i}")

            # Unpack and repack all `QuantizedTensorModule` classes in attention
            self_attn = getattr(layer, "self_attn", None) or getattr(layer, "self_attention", None)
            for _, q_tensors in self_attn.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    # Set `g_idx` to None since it's not used in `MatMulNBits`
                    q_tensors.g_idx = None

            # Unpack and repack all `QuantizedTensorModule` classes in MLP
            for _, q_tensors in layer.mlp.__dict__.items():
                if isinstance(q_tensors, QuantizedTensorModule) and q_tensors.qweight is not None:
                    self.unpack(q_tensors)
                    self.repack(q_tensors)

                    # Set `g_idx` to None since it's not used in `MatMulNBits`
                    q_tensors.g_idx = None

                if isinstance(q_tensors, QuantizedExperts) and q_tensors.num_experts > 0:
                    # Process each expert
                    self.unpack_repack_experts(q_tensors)

        if isinstance(self.lm_head, QuantizedTensorModule) and self.lm_head.qweight is not None:
            self.unpack(self.lm_head)
            self.repack(self.lm_head)

            # Set `g_idx` to None since it's not used in `MatMulNBits`
            self.lm_head.g_idx = None

    def unpack_repack_experts(self, experts):
        """
        Unpacks weights from pre-quantized Quark experts
        """
        for expert in experts.values():
            # Process gate_proj
            if expert.gate_proj.qweight is not None:
                self.unpack_qzeros(expert.gate_proj)
                self.pack_zeros_ort_format(expert.gate_proj)
                self.unpack_qweight_quark(expert.gate_proj)
                expert.gate_proj.g_idx = None

            # Process up_proj
            if expert.up_proj.qweight is not None:
                self.unpack_qzeros(expert.up_proj)
                self.pack_zeros_ort_format(expert.up_proj)
                self.unpack_qweight_quark(expert.up_proj)
                expert.up_proj.g_idx = None

            # Process fused gate_up_proj
            if expert.gate_up_proj.qweight is not None:
                self.unpack_qzeros(expert.gate_up_proj)
                self.pack_zeros_ort_format(expert.gate_up_proj)
                self.unpack_qweight_quark(expert.gate_up_proj)
                expert.gate_up_proj.g_idx = None

            # Process down_proj
            if expert.down_proj.qweight is not None:
                self.unpack_qzeros(expert.down_proj)
                self.pack_zeros_ort_format(expert.down_proj)
                self.unpack_qweight_quark(expert.down_proj)
                expert.down_proj.g_idx = None

        """
        Repacks weights from pre-quantized Quark experts
        into the format expected by the QMoE operator.
        """
        self.repack_qmoe_weights(experts)

    def repack_qmoe_weights(self, experts):
        """
        Create quantized MoE weights from pre-quantized Quark experts.
        For gate_up projection, it interleaves gate_proj and up_proj tensors
        where even indices are for gate and odd indices are for up.
        """
        has_split_gate_up = all(
            expert.gate_proj is not None and expert.gate_proj.qweight is not None and expert.up_proj is not None and expert.up_proj.qweight is not None for expert in experts.values()
        )

        if has_split_gate_up:
            self.combine_and_repack_gate_up(experts)
            self.repack_projections(experts, ['down_proj'])
        else:
            self.repack_projections(experts, ['gate_up_proj', 'down_proj'])

    def repack_projections(self, experts, projection_types):
        for proj_type in projection_types:
            qweight_list = []
            scales_list = []
            zero_points_list = []

            for expert_id in sorted(experts.keys()):
                expert = experts[expert_id]
                # Handle single projections like down_proj (or any other case)
                combined_qweight_parts = []
                combined_scales_parts = []
                combined_zp_parts = []

                proj_module = getattr(expert, proj_type)
                weights = proj_module.qweight
                weights = self.repack_qweight(weights, bits=proj_module.bits)
                combined_qweight_parts.append(weights)
                combined_scales_parts.append(proj_module.scales.T)
                combined_zp_parts.append(proj_module.qzeros)

                qweight_list.append(torch.cat(combined_qweight_parts, dim=0))
                scales_list.append(torch.cat(combined_scales_parts, dim=0))
                zero_points_list.append(torch.cat(combined_zp_parts, dim=0))

            # Stack all experts' weights and scales
            qweight = torch.stack(qweight_list, dim=0)
            scales = torch.stack(scales_list, dim=0)
            zero_points = torch.stack(zero_points_list, dim=0)

            if proj_type == 'down_proj':
                experts.fc2_weights, experts.fc2_scales, experts.fc2_zero_points = (
                    qweight,
                    scales.to(torch.float16),
                    zero_points,
                )
            else:
                experts.fc1_weights, experts.fc1_scales, experts.fc1_zero_points = (
                    qweight,
                    scales.to(torch.float16),
                    zero_points,
                )

    def combine_and_repack_gate_up(self, experts):
        gate_experts_weights = []
        up_experts_weights = []
        gate_expert_scales = []
        up_expert_scales = []
        gate_expert_zero_points = []
        up_expert_zero_points = []

        # Collect weights and scales for all experts
        for expert_id in sorted(experts.keys()):
            expert = experts[expert_id]
            gate_proj = expert.gate_proj
            up_proj = expert.up_proj

            # qweight: [inter, hidden], scales: [inter, hidden // block_size]
            gate_experts_weights.append(gate_proj.qweight)
            up_experts_weights.append(up_proj.qweight)
            gate_expert_scales.append(gate_proj.scales)
            up_expert_scales.append(up_proj.scales)
            gate_expert_zero_points.append(gate_proj.qzeros)
            up_expert_zero_points.append(up_proj.qzeros)

        # Stack experts: [experts, inter, hidden]
        gate_weights = torch.stack(gate_experts_weights, axis=0)
        up_weights = torch.stack(up_experts_weights, axis=0)

        # Concatenate along last dim, then reshape to [experts, inter*2, hidden]
        fc1 = torch.concat([gate_weights, up_weights], axis=-1).view(up_weights.shape[0], up_weights.shape[1]*2, up_weights.shape[2])

        packed_weights = [self.repack_qweight(fc1[expert_id], bits=experts[expert_id].gate_proj.bits) for expert_id in sorted(experts.keys())]
        # Stack into a 3D tensor: [num_experts, inter_size * 2, hidden_size // pack_size]
        final_fc1 = torch.stack(packed_weights, dim=0)

        # Stack scales: [experts, inter, hidden // block_size]
        gate_scales = torch.stack(gate_expert_scales, axis=0).transpose(-1, -2)  # [experts, inter, hidden // 32]
        up_scales = torch.stack(up_expert_scales, axis=0).transpose(-1, -2)  # [experts, inter, hidden // 32]
        fc1_scales = torch.concat([gate_scales, up_scales], axis=-1).view(up_scales.shape[0], up_scales.shape[1]*2, up_scales.shape[2])

        gate_zero_points = torch.stack(gate_expert_zero_points, axis=0)
        up_zero_points = torch.stack(up_expert_zero_points, axis=0)
        fc1_zero_points = torch.concat([gate_zero_points, up_zero_points], axis=1)

        experts.fc1_weights, experts.fc1_scales, experts.fc1_zero_points = final_fc1, fc1_scales.to(torch.float16), fc1_zero_points

    def repack_qweight(self, weights, bits) -> torch.Tensor:
        """
        Repacks unpacked uint8 weights (representing 4-bit values) into a packed uint8 tensor.
        This mirrors the packing logic from builder.py's _symmetric_blockwise_quantize.
        """
        if bits != 4:
            raise NotImplementedError("This repacking function is specifically for 4-bit weights.")

        quantized_flat = weights.cpu()
        original_shape = quantized_flat.shape
        quantized_uint4 = quantized_flat.to(torch.uint8)

        packed_shape = list(original_shape)
        packed_shape[-1] = (original_shape[-1] + 1) // 2
        packed_weight = torch.zeros(packed_shape, dtype=torch.uint8, device=quantized_flat.device)

        # Pack two 4-bit values per byte
        for i in range(0, quantized_uint4.shape[-1], 2):
            val1 = quantized_uint4[..., i]
            if i + 1 < quantized_uint4.shape[-1]:
                val2 = quantized_uint4[..., i + 1]
                packed_val = (val1 & 0xF) | ((val2 & 0xF) << 4)
            else:
                # Odd number of values - pack only lower 4 bits
                packed_val = val1 & 0xF
            packed_weight[..., i // 2] = packed_val

        return packed_weight

    def _load_quant_config(self, quant_attrs):
        self.global_quant_config = quant_attrs["config"]["global_quant_config"]["weight"]
        self.global_group_size = self.global_quant_config["group_size"]
        global_dtype = self.global_quant_config["dtype"]

        dtype_bits_maps = {
            "uint4": 4,
            "int4": 4,
        }

        if global_dtype not in dtype_bits_maps:
            raise ValueError(f"Unexpected dtype: {global_dtype}.")
        self.global_bits = dtype_bits_maps[global_dtype]

    def get_layer_bits(self, layer_name):
        name = layer_name.split(".")[0]
        if name in self._quant_attrs["config"]["layer_quant_config"]:
            layer_quant_config = self._quant_attrs["config"]["layer_quant_config"][name]["weight"]
            local_dtype = layer_quant_config["dtype"]

            dtype_bits_maps = {
                "uint4": 4,
                "int4": 4,
            }
            if local_dtype not in dtype_bits_maps:
                raise ValueError(f"Unexpected dtype: {local_dtype}.")
            return dtype_bits_maps[local_dtype]
        return self.global_bits

    def get_layer_group_size(self, layer_name):
        name = layer_name.split(".")[0]
        if name in self._quant_attrs["config"]["layer_quant_config"]:
            layer_quant_config = self._quant_attrs["config"]["layer_quant_config"][name]["weight"]
            return layer_quant_config["group_size"]
        return self.global_group_size

    def unpack_qweight(self, module):
        """
        Unpack `qweight` to standard format
        """
        expected_shape = (module.qweight.shape[0], module.out_features)
        transpose = module.qweight.shape != expected_shape
        module.qweight = self.unpack_on_row(module.qweight.T, module.bits, transpose)
        module.qweight = self.reverse_reorder_tensor(module.qweight.T, module.bits)
        # Padding might happen on the last dimension.
        module.qweight = module.qweight[:, : module.out_features]

    def unpack_qzeros(self, module):
        """
        Unpack `qzeros` to standard format
        """
        super().unpack_qzeros(module)
        module.qzeros = self.reverse_reorder_tensor(module.qzeros, module.bits)
        # Padding might happen on the last dimension.
        module.qzeros = module.qzeros[:, : module.out_features]

    def reverse_reorder_tensor(self, tensor, bits):
        """
        Re-arrange tensor data in a new order
        """
        compress_ratio = 32 // bits
        assert tensor.shape[-1] % compress_ratio == 0

        if bits == 4:
            order_map = [0, 2, 4, 6, 1, 3, 5, 7]
        else:
            raise NotImplementedError(f"Unpacking for {bits}-bit quantization is not currently supported.")

        order_tensor = torch.tensor(order_map, dtype=torch.int32).reshape(1, -1)
        order_tensor = order_tensor.repeat(tensor.shape[1] // compress_ratio, 1)
        order_tensor = order_tensor + torch.arange(0, tensor.shape[1], compress_ratio, dtype=torch.int32).reshape(-1, 1)
        order_tensor = order_tensor.reshape(-1)

        reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
        reverse_order_tensor = reverse_order_tensor[order_tensor]
        int_tensor = tensor[:, reverse_order_tensor]
        return int_tensor

    def unpack_qweight_quark(self, module, reorder=True, dtype="uint4"):
        """
        Unpack `qweight` to standard format and reorder for OGA.
        This is based on the packing logic from Quark's Pack_4_bits with reorder=True
        and the unpacking logic from ORT GenAI's _symmetric_blockwise_quantize.
        """
        to_unpack = module.qweight

        if to_unpack.ndim > 2:
            raise ValueError("Unpack: Only supports tensors with dimensions not greater than 2.")

        shifts = torch.tensor([0, 4, 8, 12, 16, 20, 24, 28], device=to_unpack.device)
        org_ndim = to_unpack.ndim

        if org_ndim == 1:
            to_unpack = to_unpack.unsqueeze(0)

        if to_unpack.ndim == 2:
            unpacked = (to_unpack.unsqueeze(-1) >> shifts.view(1, 1, -1)).view(to_unpack.shape[0], -1).to(torch.int8)
            if reorder:
                ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
                order_tensor = torch.arange(
                    unpacked.shape[-1],
                    dtype=torch.int32,
                    device=unpacked.device,
                )
                order_tensor = order_tensor.view(-1, 8)
                order_tensor = order_tensor[:, ORDER].view(-1)
                unpacked = unpacked[:, order_tensor]
        elif to_unpack.ndim == 0:
            unpacked = to_unpack

        unpacked &= 0b1111

        if dtype == "int4":
            mask = (unpacked & 0x08).bool()
            unpacked[mask] = unpacked[mask] | 0xF0

        if org_ndim == 1:
            unpacked = unpacked.squeeze(0)

        module.qweight = unpacked.T.contiguous()

class OliveModel(GPTQModel):
    """
    Olive quantization format:
    - qweight: (out_features, packed_in_features) uint8, packed along last dim
    - scales: (out_features, num_groups) float
    - qzeros: (out_features, packed_num_groups) uint8, packed along last dim
    """

    def _load_quant_config(self, quant_attrs):
        super()._load_quant_config(quant_attrs)
        self.overrides = quant_attrs["config"]["overrides"] or {}

    def get_layer_bits(self, layer_name):
        name = ".".join(layer_name.split(".")[:-1])
        return self.overrides.get(name, {}).get("bits", self.global_bits)

    def get_layer_group_size(self, layer_name):
        name = ".".join(layer_name.split(".")[:-1])
        return self.overrides.get(name, {}).get("group_size", self.global_group_size)

    def handle_qzeros(self, module):
        """Olive uses unsigned quantization, no offset needed."""
        pass

    def unpack(self, module):
        """Skip unpack for Olive format."""
        pass

    def repack(self, module):
        """
        Olive format:
        - qweight: (out_features, packed_in_features) uint8
        - scales: (out_features, num_groups) float
        - qzeros: (out_features, packed_num_groups) uint8

        ORT format:
        - qweight: (out_features, k_blocks, blob_size) uint8
        - scales: (out_features * num_groups,) float, flattened
        - qzeros: (out_features * packed_num_groups,) uint8, flattened
        """
        kpack = 8 // module.bits
        k_blocks = module.in_features // module.group_size
        blob_size = module.group_size // kpack

        # qweight: (out_features, packed_in_features) -> (out_features, k_blocks, blob_size)
        module.qweight = module.qweight.reshape(module.out_features, k_blocks, blob_size).contiguous()

        # scales: (out_features, num_groups) -> flatten to 1D
        module.scales = module.scales.reshape(-1).contiguous()

        # qzeros: (out_features, packed_num_groups) -> flatten to 1D
        if module.qzeros is not None and module.qzeros.numel() > 0:
            module.qzeros = module.qzeros.reshape(-1).contiguous()


# ---------------------------------------------------------------------------
# NVIDIA Model Optimizer (modelopt) NVFP4 + FP8 checkpoints
# ---------------------------------------------------------------------------
# Model Optimizer exports a mixed-precision HF checkpoint:
#   * MoE experts + shared expert + lm_head : W4A16_NVFP4 (block-16 E2M1 weights,
#     FP8-E4M3 block scales, per-tensor FP32 global scale `weight_scale_2`).
#   * attention (self_attn / linear_attn projections) : FP8 (per-tensor weight_scale).
#   * everything else (norms, conv1d, router, embeddings) : BF16.
#
# The ONNX Runtime GenAI QMoE op consumes the NVFP4 *routed* experts natively
# (quant_type="nvfp4"), and the Qwen builder reads those per-expert tensors
# straight from the source safetensors (see make_nvfp4_moe_initializers). So this
# loader only has to materialize the NON-routed weights the builder walks as plain
# modules: it dequantizes the FP8 attention and NVFP4 shared-expert / lm_head to
# BF16 and passes BF16 tensors through unchanged. Attention is exported through the
# builder's normal (int4 RTN) path since ONNX Runtime has no native FP8 attention
# GEMM; dequantizing the FP8 weights to BF16 reconstructs them exactly.


class ModeloptDecoderLayer:
    """Lightweight decoder-layer container (name ends in 'DecoderLayer' so the
    builder's is_layer() recognizes it). Only the attention variant present in the
    checkpoint (linear_attn or self_attn) is populated; the other stays None."""

    def __init__(self, layer_id):
        self.layer_id = layer_id
        self.input_layernorm = TensorModule()
        self.post_attention_layernorm = TensorModule()
        self.self_attn = None
        self.linear_attn = None
        self.mlp = None

    def is_empty(self):
        return self.input_layernorm.weight is None


class ModeloptModel(QuantizedModel):
    """Loader for NVIDIA Model Optimizer NVFP4 + FP8 mixed-precision checkpoints.

    The base initializes the module surface the builder walks without running its
    eager integer-checkpoint loading path. This loader instead streams tensors on
    demand and dequantizes only the non-routed weights to BF16.
    """

    _FP4_E2M1_LUT = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32)

    def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers):
        import json
        from types import SimpleNamespace
        from safetensors import safe_open

        with open(os.path.join(input_path, "config.json")) as f:
            cfg = json.load(f)
        text_config = cfg.get("text_config", cfg)
        num_layers = num_layers or text_config.get("num_hidden_layers")
        super().__init__(
            quant_type, input_path, quant_attrs, q_size, kv_size, intermediate_size, num_layers, load_weights=False
        )
        self._input_path = input_path
        self._safe_open = safe_open
        self._simple_namespace = SimpleNamespace
        self._open_handles = {}
        self._handle_keys = {}

        index_path = os.path.join(input_path, "model.safetensors.index.json")
        if os.path.exists(index_path):
            with open(index_path) as f:
                self._weight_map = json.load(f)["weight_map"]
            self._single_file = None
        else:
            self._weight_map = None
            candidates = sorted(f for f in os.listdir(input_path) if f.endswith(".safetensors"))
            if len(candidates) != 1:
                raise ValueError(
                    f"'{input_path}' has no 'model.safetensors.index.json', so it must contain exactly one "
                    f".safetensors file, but found {len(candidates)}: {candidates}."
                )
            self._single_file = candidates[0]

        try:
            self.layers.extend(self._build_layer(layer_id) for layer_id in range(num_layers))

            # Globals: embeddings + final norm are BF16; lm_head is NVFP4.
            self.embedding.weight = self._get("model.language_model.embed_tokens.weight")
            self.final_norm.weight = self._get("model.language_model.norm.weight")
            self.lm_head.weight = self._dequant_linear("lm_head")
        finally:
            # Every tensor is materialized above; do not hold file descriptors open for
            # the rest of the (long) export. `_get` re-opens lazily if it is called again.
            self.close()

    def close(self):
        """Release the cached safetensors file handles."""
        for handle in self._open_handles.values():
            handle.__exit__(None, None, None)
        self._open_handles.clear()
        self._handle_keys.clear()

    def __del__(self):
        if getattr(self, "_open_handles", None):
            self.close()

    def repack_nvfp4_weight_codes(self, packed_nk2):
        """Unpack a Model Optimizer NVFP4 weight tensor to per-element e2m1 codes."""
        if packed_nk2.dtype != torch.uint8:
            packed_nk2 = packed_nk2.to(torch.uint8)
        low = packed_nk2 & 0x0F
        high = packed_nk2 >> 4
        n = packed_nk2.shape[0]
        return torch.stack((low, high), dim=-1).reshape(n, -1).contiguous()

    def pack_nvfp4_codes_for_qmoe(self, codes_nk):
        """Pack e2m1 codes ``[N, K]`` into the CUDA QMoE ``[K, N/2]`` layout."""
        if codes_nk.dtype != torch.uint8:
            codes_nk = codes_nk.to(torch.uint8)
        n = codes_nk.shape[0]
        if n % 2 != 0:
            raise ValueError(f"NVFP4 QMoE packing requires an even N={n} for nibble packing.")
        codes_kn = codes_nk.T.contiguous()
        low = codes_kn[:, 0::2] & 0x0F
        high = codes_kn[:, 1::2] & 0x0F
        return ((high << 4) | low).contiguous()

    def _dequant_nvfp4(self, weight_u8, block_scale_e4m3, global_scale, name=""):
        """Reconstruct a BF16 weight from Model Optimizer NVFP4 tensors."""
        if block_scale_e4m3 is None:
            raise ValueError(
                f"NVFP4 tensor '{name}' has 'weight_scale_2' but no 'weight_scale' (FP8-E4M3 block scales). "
                "The Model Optimizer checkpoint is incomplete."
            )
        if block_scale_e4m3.dtype == torch.uint8:
            block_scale_e4m3 = block_scale_e4m3.view(torch.float8_e4m3fn)
        elif block_scale_e4m3.dtype != torch.float8_e4m3fn:
            raise ValueError(
                f"NVFP4 tensor '{name}' block scales must be float8_e4m3fn (or raw uint8 bytes), "
                f"got {block_scale_e4m3.dtype}."
            )
        codes = self.repack_nvfp4_weight_codes(weight_u8).long()
        mag = self._FP4_E2M1_LUT[codes & 0x7]
        values = torch.where((codes & 0x8) > 0, -mag, mag)
        block_scales = block_scale_e4m3.to(torch.float32)
        n, k = codes.shape
        if block_scales.shape[0] != n or block_scales.shape[1] == 0 or k % block_scales.shape[1] != 0:
            raise ValueError(
                f"NVFP4 tensor '{name}' block scales {tuple(block_scales.shape)} are not a block-wise split of the "
                f"[{n}, {k}] weight."
            )
        block_scales = block_scales.repeat_interleave(k // block_scales.shape[1], dim=1)
        return (values * block_scales * float(global_scale)).to(torch.bfloat16)

    def _dequant_fp8(self, weight_f8, weight_scale):
        """Reconstruct a BF16 weight from an FP8 (E4M3) weight and per-tensor scale."""
        return (weight_f8.to(torch.float32) * float(weight_scale)).to(torch.bfloat16)

    # -- raw tensor access ------------------------------------------------------
    def _get(self, name):
        if self._weight_map is not None:
            fname = self._weight_map.get(name)
            files = [fname] if fname is not None else []
        else:
            files = [self._single_file]
        for fname in files:
            handle = self._open_handles.get(fname)
            if handle is None:
                handle = self._safe_open(os.path.join(self._input_path, fname), framework="pt", device="cpu")
                self._open_handles[fname] = handle
                self._handle_keys[fname] = set(handle.keys())
            if name in self._handle_keys[fname]:
                return handle.get_tensor(name)
        return None

    # -- dequantized linear weight ---------------------------------------------
    def _dequant_linear(self, base):
        weight = self._get(f"{base}.weight")
        if weight is None:
            return None
        weight_scale_2 = self._get(f"{base}.weight_scale_2")
        weight_scale = self._get(f"{base}.weight_scale")
        if weight_scale_2 is not None:  # NVFP4 (block-16 E2M1 + E4M3 block scale + global)
            return self._dequant_nvfp4(weight, weight_scale, weight_scale_2, name=base)
        if weight_scale is not None and weight.dtype == torch.float8_e4m3fn:  # FP8 (per-tensor scale)
            return self._dequant_fp8(weight, weight_scale)
        return weight.to(torch.bfloat16)

    def _linear_module(self, base):
        weight = self._dequant_linear(base)
        if weight is None:
            return None
        module = TensorModule()
        module.weight = weight
        bias = self._get(f"{base}.bias")
        if bias is not None:
            module.bias = bias
        return module

    def _tensor_module(self, name):
        module = TensorModule()
        module.weight = self._get(name)
        return module

    def _build_layer(self, layer_id):
        ns = self._simple_namespace
        p = f"model.language_model.layers.{layer_id}"
        layer = ModeloptDecoderLayer(layer_id)
        layer.input_layernorm.weight = self._get(f"{p}.input_layernorm.weight")
        layer.post_attention_layernorm.weight = self._get(f"{p}.post_attention_layernorm.weight")

        if self._get(f"{p}.linear_attn.in_proj_qkv.weight") is not None:
            la = ns()
            la.in_proj_qkv = self._linear_module(f"{p}.linear_attn.in_proj_qkv")
            la.in_proj_z = self._linear_module(f"{p}.linear_attn.in_proj_z")
            la.in_proj_a = self._linear_module(f"{p}.linear_attn.in_proj_a")
            la.in_proj_b = self._linear_module(f"{p}.linear_attn.in_proj_b")
            la.out_proj = self._linear_module(f"{p}.linear_attn.out_proj")
            la.conv1d = self._tensor_module(f"{p}.linear_attn.conv1d.weight")
            la.A_log = self._get(f"{p}.linear_attn.A_log")
            la.dt_bias = self._get(f"{p}.linear_attn.dt_bias")
            la.norm = self._tensor_module(f"{p}.linear_attn.norm.weight")
            layer.linear_attn = la
        elif self._get(f"{p}.self_attn.q_proj.weight") is not None:
            sa = ns()
            sa.q_proj = self._linear_module(f"{p}.self_attn.q_proj")
            sa.k_proj = self._linear_module(f"{p}.self_attn.k_proj")
            sa.v_proj = self._linear_module(f"{p}.self_attn.v_proj")
            sa.o_proj = self._linear_module(f"{p}.self_attn.o_proj")
            sa.q_norm = self._tensor_module(f"{p}.self_attn.q_norm.weight")
            sa.k_norm = self._tensor_module(f"{p}.self_attn.k_norm.weight")
            layer.self_attn = sa
        else:
            raise ValueError(
                f"Layer {layer_id} has neither '{p}.linear_attn.in_proj_qkv.weight' nor "
                f"'{p}.self_attn.q_proj.weight', so its attention variant cannot be determined. "
                "The Model Optimizer checkpoint is incomplete or uses unsupported weight names."
            )

        mlp = ns()
        mlp.gate = self._tensor_module(f"{p}.mlp.gate.weight")
        shared = ns()
        shared.gate_proj = self._linear_module(f"{p}.mlp.shared_expert.gate_proj")
        shared.up_proj = self._linear_module(f"{p}.mlp.shared_expert.up_proj")
        shared.down_proj = self._linear_module(f"{p}.mlp.shared_expert.down_proj")
        mlp.shared_expert = shared
        mlp.shared_expert_gate = self._tensor_module(f"{p}.mlp.shared_expert_gate.weight")
        # Routed experts are streamed from raw safetensors by the builder's
        # make_nvfp4_moe_initializers(); nothing to materialize here.
        mlp.experts = None
        layer.mlp = mlp
        return layer


class QuantModel:
    @staticmethod
    def from_pretrained(quant_type, **kwargs):
        """
        Unpack quantized weights in PyTorch models, store them in a standard format, and repack them
        into ONNX Runtime's format. Also performs any pre-processing and post-processing when unpacking
        the quantized weights.
        """
        if quant_type == "awq":
            model = AWQModel(quant_type, **kwargs)
        elif quant_type == "gptq":
            model = GPTQModel(quant_type, **kwargs)
        elif quant_type == "olive":
            model = OliveModel(quant_type, **kwargs)
        elif quant_type == "quark":
            model = QuarkModel(quant_type, **kwargs)
        elif quant_type == "modelopt":
            model = ModeloptModel(quant_type, **kwargs)
        else:
            raise NotImplementedError(f"The {quant_type} quantized model is not currently supported.")

        return model
