chore: import upstream snapshot with attribution

python/mlc_llm/quantization/__init__.py

@@ -0,0 +1,11 @@
+"""A subpackage for quantization and dequantization algorithms"""
+
+from .awq_quantization import AWQQuantize
+from .block_scale_quantization import BlockScaleQuantize
+from .fp8_quantization import FP8PerTensorQuantizeMixtralExperts
+from .ft_quantization import FTQuantize
+from .group_quantization import GroupQuantize
+from .model_quantization import make_awq_quant, make_quantization_functions
+from .no_quantization import NoQuantize
+from .per_tensor_quantization import PerTensorQuantize
+from .quantization import QUANTIZATION, Quantization

python/mlc_llm/quantization/awq_quantization.py

@@ -0,0 +1,282 @@
+"""AWQ Quantization"""
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional  # noqa: UP035
+
+from tvm import DataType, DataTypeCode, te, tirx, topi
+from tvm.relax.frontend import nn
+from tvm.runtime import Tensor
+
+from mlc_llm.loader import QuantizeMapping
+
+from .utils import convert_uint_to_float, is_final_fc, is_moe_gate
+
+
+def _make_divisible(c, divisor):
+    return (c + divisor - 1) // divisor
+
+
+def _calculate_zeros_width(in_features, group_size=128, pack_num=8):
+    if group_size >= 128:
+        size_multiplier = 1
+    elif group_size == 64:
+        size_multiplier = 2
+    elif group_size == 32:
+        size_multiplier = 4
+    else:
+        raise NotImplementedError
+
+    base_width = _make_divisible(in_features // group_size, pack_num)
+    base_width = _make_divisible(base_width, size_multiplier) * size_multiplier
+    return base_width
+
+
+@dataclass
+class AWQQuantize:
+    """Configuration for AWQ quantization"""
+
+    name: str
+    kind: str
+    group_size: int
+    quantize_dtype: str  # "int3", "int4", "int8"
+    storage_dtype: str  # "uint32"
+    model_dtype: str  # "float16", "float32"
+
+    num_elem_per_storage: int = 0
+    num_storage_per_group: int = 0
+    max_int_value: int = 0
+
+    prebuilt_quantize_func: Dict[str, Callable[[Tensor], Tensor]] = field(  # noqa: UP006
+        default_factory=lambda: {}
+    )
+
+    def __post_init__(self):
+        assert self.kind == "awq"
+        quantize_dtype = DataType(self.quantize_dtype)
+        storage_dtype = DataType(self.storage_dtype)
+        model_dtype = DataType(self.model_dtype)
+        assert quantize_dtype.type_code == DataTypeCode.INT
+        assert storage_dtype.type_code == DataTypeCode.UINT
+        assert model_dtype.type_code == DataTypeCode.FLOAT
+        if storage_dtype.bits < quantize_dtype.bits:
+            raise ValueError("Storage unit should be greater or equal to quantized element")
+
+        self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits
+        if self.group_size % self.num_elem_per_storage != 0:
+            raise ValueError("Group size should be divisible by numbers of elements per storage")
+        self.num_storage_per_group = self.group_size // self.num_elem_per_storage
+        self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1
+
+    def quantize_model(
+        self,
+        model: nn.Module,
+        quant_map: QuantizeMapping,
+        name_prefix: str,
+    ) -> nn.Module:
+        """
+        Quantize model with awq quantization.
+
+        Parameters
+        ----------
+        model : nn.Module
+            The non-quantized nn.Module.
+
+        quant_map : QuantizeMapping
+            The quantize mapping with name mapping and func mapping.
+
+        name_prefix : str
+            The name prefix for visited weight.
+
+        Returns
+        -------
+        ret : nn.Module
+            The quantized nn.Module.
+        """
+
+        class _Mutator(nn.Mutator):
+            def __init__(self, config: AWQQuantize, quant_map: QuantizeMapping) -> None:
+                super().__init__()
+                self.config = config
+                self.quant_map = quant_map
+
+            def visit_module(self, name: str, node: nn.Module) -> Any:
+                """
+                The visiting method for awq quantization of nn.Module nodes.
+
+                Parameters
+                ----------
+                name : str
+                    The name of the current node
+
+                node : nn.Module
+                    The current node of nn.Module to mutate.
+
+                Returns
+                -------
+                ret_node : Any
+                    The new node to replace current node.
+                """
+
+                if (
+                    isinstance(node, nn.Linear)
+                    and not is_final_fc(name)
+                    and not is_moe_gate(name, node)
+                ):
+                    return AWQQuantizeLinear.from_linear(node, self.config)
+                return self.visit(name, node)
+
+        model.to(dtype=self.model_dtype)
+        mutator = _Mutator(self, quant_map)
+        model = mutator.visit(name_prefix, model)
+        return model
+
+    def _dequantize(
+        self,
+        weight: te.Tensor,
+        zeros: te.Tensor,
+        scale: te.Tensor,
+        out_shape: Optional[List[tirx.Expr]] = None,  # noqa: UP006
+    ):
+        float_weight = convert_uint_to_float(
+            weight,
+            DataType(self.quantize_dtype).bits,
+            self.num_elem_per_storage,
+            self.storage_dtype,
+            self.model_dtype,
+            out_shape=[weight.shape[0], weight.shape[1] * self.num_elem_per_storage],
+            ft_reorder=True,
+        )
+        float_zeros = convert_uint_to_float(
+            zeros,
+            DataType(self.quantize_dtype).bits,
+            self.num_elem_per_storage,
+            self.storage_dtype,
+            self.model_dtype,
+            out_shape=[zeros.shape[0], zeros.shape[1] * self.num_elem_per_storage],
+            ft_reorder=True,
+        )
+        float_weight = topi.transpose(float_weight)
+        float_zeros = topi.transpose(float_zeros)
+        scale = topi.transpose(scale)
+        return te.compute(
+            shape=(
+                [weight.shape[0], weight.shape[1] * self.num_elem_per_storage]
+                if out_shape is None
+                else out_shape
+            ),
+            fcompute=lambda i, j: tirx.Mul(
+                tirx.Sub(float_weight[i, j], float_zeros[i, j // self.group_size]),
+                scale[i, j // self.group_size],
+            ),
+            name="dequantize",
+        )
+
+
+class AWQQuantizeLinear(nn.Module):
+    """An nn.Linear module with AWQ quantization"""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        config: AWQQuantize,
+        bias: bool = True,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype
+        self.config = config
+        self.qweight = nn.Parameter(
+            (in_features, out_features // config.num_elem_per_storage),
+            config.storage_dtype,
+        )
+        self.qzeros = nn.Parameter(
+            (
+                in_features // config.group_size,
+                out_features // config.num_elem_per_storage,
+            ),
+            config.storage_dtype,
+        )
+        self.scales = nn.Parameter(
+            (in_features // config.group_size, out_features), config.model_dtype
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                (out_features,), config.model_dtype if out_dtype is None else out_dtype
+            )
+        else:
+            self.bias = None
+
+    @staticmethod
+    def from_linear(linear: nn.Linear, config: AWQQuantize) -> "AWQQuantizeLinear":
+        """
+        Converts a non-quantized nn.Linear to a group quantized AWQQuantizeLinear
+
+        Parameters
+        ----------
+        linear : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : AWQQuantize
+            The awq quantization config.
+
+        Returns
+        -------
+        ret : GroupQuantizeLinear
+            The awq quantized AWQQuantizeLinear layer.
+        """
+        return AWQQuantizeLinear(
+            in_features=linear.in_features,
+            out_features=linear.out_features,
+            config=config,
+            bias=getattr(linear, "bias", None) is not None,
+            out_dtype=linear.out_dtype,
+        )
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        """
+        Forward method for awq quantized linear layer
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the group quantized linear layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, zeros, scale: self.config._dequantize(
+                weight,
+                zeros,
+                scale,
+                [
+                    tirx.IntImm("int64", self.out_features),
+                    tirx.IntImm("int64", self.in_features),
+                ],
+            ),
+            name_hint="dequantize",
+            args=[self.qweight, self.qzeros, self.scales],
+        )
+        w = nn.op.permute_dims(w)
+        x = nn.op.matmul(x, w, out_dtype=self.out_dtype)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        self.qweight.to(dtype=dtype)
+        self.qzeros.to(dtype=dtype)
+        self.scales.to(dtype=dtype)
+        if self.bias is not None and self.out_dtype is None:
+            self.bias.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype

python/mlc_llm/quantization/block_scale_quantization.py

@@ -0,0 +1,828 @@
+"""The block-scale quantization config"""
+
+from dataclasses import dataclass
+from typing import Any, Literal, Optional, Tuple  # noqa: UP035
+
+import tvm
+from tvm import DataType, DataTypeCode, te, tirx
+from tvm.relax.frontend import nn
+from tvm.script import tirx as T
+
+from mlc_llm.loader import QuantizeMapping
+from mlc_llm.nn import MixtralExperts
+from mlc_llm.op import cutlass, extern, moe_matmul, triton
+from mlc_llm.support import logging
+from mlc_llm.support import tensor_parallel as tp
+
+from .utils import apply_sharding, is_final_fc, is_moe_gate
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BlockScaleQuantize:
+    """Configuration for block-scale quantization"""
+
+    name: str
+    kind: str = "block-scale"
+    weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"] = "float8_e4m3fn"
+    model_dtype: Literal["float16", "bfloat16"] = "bfloat16"
+    quantize_linear: bool = True
+    weight_block_size: Optional[Tuple[int, int]] = None  # noqa: UP006
+    use_activation_scale: bool = False
+
+    def __post_init__(self):
+        assert self.kind == "block-scale-quant"
+        weight_dtype = DataType(self.weight_dtype)
+        model_dtype = DataType(self.model_dtype)
+        assert weight_dtype.type_code in [
+            DataTypeCode.Float8E4M3FN,
+            DataTypeCode.Float8E5M2,
+        ]
+        assert model_dtype.type_code in [
+            DataTypeCode.FLOAT,
+            DataTypeCode.BFLOAT,
+        ]
+
+    def quantize_model(
+        self,
+        model: nn.Module,
+        quant_map: QuantizeMapping,
+        name_prefix: str,
+    ) -> nn.Module:
+        """Quantize model with block-scale quantization
+
+        Parameters
+        ----------
+        model : nn.Module
+            The non-quantized nn.Module.
+
+        quant_map : QuantizeMapping
+            The quantize mapping with name mapping and func mapping.
+
+        name_prefix : str
+            The name prefix for visited weight.
+
+        Returns
+        -------
+        ret : nn.Module
+            The quantized nn.Module.
+        """
+
+        weight_block_size = model.weight_block_size
+
+        class _Mutator(nn.Mutator):
+            def __init__(self, config: BlockScaleQuantize, quant_map: QuantizeMapping) -> None:
+                super().__init__()
+                self.config = config
+                self.quant_map = quant_map
+
+            def visit_module(self, name: str, node: nn.Module) -> Any:
+                """The visiting method for block-scale quantization of nn.Module nodes.
+
+                Parameters
+                ----------
+                name : str
+                    The name of the current node.
+
+                node : nn.Module
+                    The current node of nn.Module to mutate.
+
+                Returns
+                ------
+                ret : Any
+                """
+                if getattr(node, "no_quantization", False):
+                    return node
+
+                if hasattr(node, "w_uk"):
+                    assert hasattr(node, "w_uv")
+                    assert node.block_size == weight_block_size
+                    if (
+                        node.qk_nope_head_dim % node.block_size[0] != 0
+                        or node.v_head_dim % node.block_size[1] != 0
+                    ):
+                        raise ValueError(
+                            "Invalid DeepSeek model config: "
+                            "qk_nope_head_dim must be multiple of weight_block_size[0], and "
+                            "v_head_dim must be multiple of weight_block_size[1]. "
+                            f"However, qk_nope_head_dim is {node.qk_nope_head_dim}, "
+                            f"v_head_dim is {node.v_head_dim}, "
+                            f"weight_block_size is {node.block_size}."
+                        )
+                    w_uk_shard_strategy = node.w_uk.attrs.get("shard_strategy", None)
+                    w_uv_shard_strategy = node.w_uv.attrs.get("shard_strategy", None)
+                    node.w_uk = nn.Parameter(
+                        (node.num_heads, node.kv_lora_rank, node.qk_nope_head_dim),
+                        self.config.weight_dtype,
+                    )
+                    node.w_uv = nn.Parameter(
+                        (node.num_heads, node.v_head_dim, node.kv_lora_rank),
+                        self.config.weight_dtype,
+                    )
+                    node.w_uk_scale_inv = nn.Parameter(
+                        (
+                            node.num_heads,
+                            node.kv_lora_rank // node.block_size[1],
+                            node.qk_nope_head_dim // node.block_size[0],
+                        ),
+                        "float32",
+                    )
+                    node.w_uv_scale_inv = nn.Parameter(
+                        (
+                            node.num_heads,
+                            node.v_head_dim // node.block_size[0],
+                            node.kv_lora_rank // node.block_size[1],
+                        ),
+                        "float32",
+                    )
+                    if w_uk_shard_strategy is not None:
+                        assert w_uk_shard_strategy.segs is None
+                        apply_sharding(w_uk_shard_strategy, w_uk_shard_strategy.name, node.w_uk)
+                        apply_sharding(
+                            w_uk_shard_strategy,
+                            f"{w_uk_shard_strategy.name}_scale_inv",
+                            node.w_uk_scale_inv,
+                        )
+                    if w_uv_shard_strategy is not None:
+                        assert w_uv_shard_strategy.segs is None
+                        apply_sharding(w_uv_shard_strategy, w_uv_shard_strategy.name, node.w_uv)
+                        apply_sharding(
+                            w_uv_shard_strategy,
+                            f"{w_uv_shard_strategy.name}_scale_inv",
+                            node.w_uv_scale_inv,
+                        )
+
+                if (
+                    isinstance(node, nn.Linear)
+                    and not is_final_fc(name)
+                    and not is_moe_gate(name, node)
+                ):
+                    if self.config.use_activation_scale:
+                        return BlockScaleQuantizeLinearStaticActivation.from_linear(
+                            node, self.config, weight_block_size
+                        )
+                    return BlockScaleQuantizeLinear.from_linear(
+                        node, self.config, weight_block_size
+                    )
+                if isinstance(node, MixtralExperts):
+                    return BlockScaleQuantizeMixtralExperts.from_mixtral_experts(
+                        node, self.config, weight_block_size
+                    )
+                return self.visit(name, node)
+
+        model.to(dtype=self.model_dtype)
+        mutator = _Mutator(self, quant_map)
+        model = mutator.visit(name_prefix, model)
+        self.weight_block_size = weight_block_size
+        return model
+
+
+class BlockScaleQuantizeLinear(nn.Module):
+    """Block-scale quantization for Linear"""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"],
+        block_size: Tuple[int, int],  # noqa: UP006
+        bias: bool = True,
+        dtype: Optional[str] = None,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype
+        self.weight = nn.Parameter((out_features, in_features), weight_dtype)
+        self.weight_scale_inv = nn.Parameter(
+            (
+                (out_features + block_size[0] - 1) // block_size[0],
+                (in_features + block_size[1] - 1) // block_size[1],
+            ),
+            "float32",
+        )
+        self.weight_dtype = weight_dtype
+        self.block_size = block_size
+        if bias:
+            self.bias = nn.Parameter((out_features,), dtype if out_dtype is None else out_dtype)
+        else:
+            self.bias = None
+
+    @staticmethod
+    def from_linear(
+        src: nn.Linear,
+        config: BlockScaleQuantize,
+        weight_block_size: Optional[Tuple[int, int]],  # noqa: UP006
+    ) -> "BlockScaleQuantizeLinear":
+        """
+        Converts a non-quantized nn.Linear to a block-scale quantized BlockScaleQuantizeLinear
+
+        Parameters
+        ----------
+        src : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : BlockScaleQuantize
+            The block-scale quantization config.
+
+        weight_block_size : Optional[Tuple[int, int]]
+            The weight block size.
+
+        Returns
+        -------
+        ret : BlockScaleQuantizeLinear
+            The block-scale quantized BlockScaleQuantizeLinear.
+        """
+        assert weight_block_size is not None
+        out_features, in_features = src.weight.shape
+        quantized_linear = BlockScaleQuantizeLinear(
+            in_features=in_features,
+            out_features=out_features,
+            weight_dtype=config.weight_dtype,
+            block_size=weight_block_size,
+            bias=getattr(src, "bias", None) is not None,
+            dtype=config.model_dtype,
+            out_dtype=src.out_dtype,
+        )
+        if quantized_linear.bias is not None:
+            quantized_linear.bias.attrs = src.bias.attrs
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, shard.name, quantized_linear.weight)
+            if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None:
+                shard.segs = [x // weight_block_size[shard.dim] for x in shard.segs]
+            apply_sharding(shard, f"{shard.name}_scale_inv", quantized_linear.weight_scale_inv)
+        return quantized_linear
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        """Forward pass of the block-scale quantized linear layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor.
+        """
+        m = 1
+        for i in range(x.ndim - 1):
+            m *= x.shape[i]
+        if m == 1:
+            x_shape = x.shape
+            return dequantize_float8_groupwise_scaled_gemv(
+                x.reshape(1, x.shape[-1]),
+                self.weight,
+                self.weight_scale_inv,
+                self.block_size,
+                self.out_dtype if self.out_dtype is not None else x.dtype,
+            ).reshape(*x_shape[:-1], -1)
+
+        shape_supported_by_cutlass = (  # noqa: F841
+            self.weight.shape[0] % 128 == 0 and self.weight.shape[1] % 128 == 0
+        )
+        # Todo: check "shape supported by cutlass" for Hopper
+        if (
+            extern.get_store().cutlass_gemm
+            and tvm.get_global_func(
+                "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn", allow_missing=True
+            )
+            is not None
+        ):
+            x_fp8, x_scale = rowwise_group_quant_fp8(
+                x, self.block_size[1], self.weight_dtype, transpose_scale=True
+            )
+            x = cutlass.fp8_groupwise_scaled_gemm(
+                x_fp8,
+                x_scale,
+                self.weight,
+                self.weight_scale_inv,
+                self.block_size,
+                self.out_dtype if self.out_dtype is not None else x.dtype,
+            )
+        else:
+            x_fp8, x_scale = rowwise_group_quant_fp8(
+                x, self.block_size[1], self.weight_dtype, transpose_scale=False
+            )
+            x = triton.fp8_groupwise_scaled_gemm(
+                x_fp8,
+                x_scale,
+                self.weight,
+                self.weight_scale_inv,
+                self.block_size,
+                self.out_dtype if self.out_dtype is not None else x.dtype,
+            )
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        if self.bias is not None and self.out_dtype is None:
+            self.bias.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype
+
+
+class BlockScaleQuantizeLinearStaticActivation(BlockScaleQuantizeLinear):
+    """Block-scale quantization for static activation FP8."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"],
+        block_size: Tuple[int, int],  # noqa: UP006
+        bias: bool = True,
+        dtype: Optional[str] = None,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__(
+            in_features=in_features,
+            out_features=out_features,
+            weight_dtype=weight_dtype,
+            block_size=block_size,
+            bias=bias,
+            dtype=dtype,
+            out_dtype=out_dtype,
+        )
+        num_in_groups = (in_features + block_size[1] - 1) // block_size[1]
+        self.activation_scale = nn.Parameter((num_in_groups,), "float32")
+
+    @staticmethod
+    def from_linear(
+        src: nn.Linear,
+        config: BlockScaleQuantize,
+        weight_block_size: Optional[Tuple[int, int]],  # noqa: UP006
+    ) -> "BlockScaleQuantizeLinearStaticActivation":
+        """
+        Convert a non-quantized nn.Linear to a block-scale quantized BlockScaleQuantizeLinearStaticActivation.
+
+        Parameters
+        ----------
+        src : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : BlockScaleQuantize
+            The block-scale quantization config.
+
+        weight_block_size : Optional[Tuple[int, int]]
+            The weight block size.
+
+        Returns
+        -------
+        ret : BlockScaleQuantizeLinearStaticActivation
+            The block-scale quantized BlockScaleQuantizeLinearStaticActivation
+        """  # noqa: E501
+        assert weight_block_size is not None
+        out_features, in_features = src.weight.shape
+        quantized_linear = BlockScaleQuantizeLinearStaticActivation(
+            in_features=in_features,
+            out_features=out_features,
+            weight_dtype=config.weight_dtype,
+            block_size=weight_block_size,
+            bias=getattr(src, "bias", None) is not None,
+            dtype=config.model_dtype,
+            out_dtype=src.out_dtype,
+        )
+        if quantized_linear.bias is not None:
+            quantized_linear.bias.attrs = src.bias.attrs
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, shard.name, quantized_linear.weight)
+            if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None:
+                shard.segs = [x // weight_block_size[shard.dim] for x in shard.segs]
+            apply_sharding(shard, f"{shard.name}_scale_inv", quantized_linear.weight_scale_inv)
+            apply_sharding(
+                shard,
+                f"{shard.name}_activation_scale",
+                quantized_linear.activation_scale,
+            )
+        return quantized_linear
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        x_fp8 = static_activation_group_quant_fp8(
+            x,
+            self.activation_scale,
+            self.block_size[1],
+            self.weight_dtype,
+        )
+        shape_supported_by_cutlass = (
+            self.weight.shape[0] % 128 == 0 and self.weight.shape[1] % 128 == 0
+        )
+        if (
+            extern.get_store().cutlass_gemm
+            and shape_supported_by_cutlass
+            and tvm.get_global_func(
+                "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn", allow_missing=True
+            )
+            is not None
+        ):
+            x_scale = broadcast_activation_scale(
+                x,
+                self.activation_scale,
+                transpose=True,
+            )
+            out = cutlass.fp8_groupwise_scaled_gemm(
+                x_fp8,
+                x_scale,
+                self.weight,
+                self.weight_scale_inv,
+                self.block_size,
+                self.out_dtype if self.out_dtype is not None else x.dtype,
+            )
+        else:
+            x_scale_triton = broadcast_activation_scale(
+                x,
+                self.activation_scale,
+                transpose=False,
+            )
+            out = triton.fp8_groupwise_scaled_gemm(
+                x_fp8,
+                x_scale_triton,
+                self.weight,
+                self.weight_scale_inv,
+                self.block_size,
+                self.out_dtype if self.out_dtype is not None else x.dtype,
+            )
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+class BlockScaleQuantizeMixtralExperts(nn.Module):
+    """Block-scale quantization for MoE experts"""
+
+    def __init__(
+        self,
+        num_local_experts: int,
+        in_features: int,
+        out_features: int,
+        weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"],
+        block_size: Tuple[int, int],  # noqa: UP006
+    ) -> None:
+        super().__init__()
+        self.num_local_experts = num_local_experts
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter((num_local_experts, out_features, in_features), weight_dtype)
+        self.weight_scale_inv = nn.Parameter(
+            (
+                num_local_experts,
+                (out_features + block_size[0] - 1) // block_size[0],
+                (in_features + block_size[1] - 1) // block_size[1],
+            ),
+            "float32",
+        )
+        self.weight_dtype = weight_dtype
+        self.block_size = block_size
+
+    @staticmethod
+    def from_mixtral_experts(
+        src: "MixtralExperts",
+        config: BlockScaleQuantize,
+        weight_block_size: Optional[Tuple[int, int]],  # noqa: UP006
+    ) -> "BlockScaleQuantizeMixtralExperts":
+        """
+        Converts a non-quantized MixtralExperts to a block-scale
+        quantized BlockScaleQuantizeMixtralExperts
+
+        Parameters
+        ----------
+        src : MixtralExperts
+            The non-quantized MixtralExperts
+
+        config : BlockScaleQuantize
+            The block-scale quantization config.
+
+        weight_block_size : Optional[Tuple[int, int]]
+            The weight block size.
+
+        Returns
+        -------
+        ret : BlockScaleQuantizeMixtralExperts
+            The block-scale quantized BlockScaleQuantizeMixtralExperts layer.
+        """
+        assert weight_block_size is not None
+        quantized_mistral_experts = BlockScaleQuantizeMixtralExperts(
+            num_local_experts=src.num_local_experts,
+            in_features=src.in_features,
+            out_features=src.out_features,
+            weight_dtype=config.weight_dtype,
+            block_size=weight_block_size,
+        )
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, shard.name, quantized_mistral_experts.weight)
+            if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None:
+                shard.segs = [x // weight_block_size[shard.dim - 1] for x in shard.segs]
+            apply_sharding(
+                shard,
+                f"{shard.name}_scale_inv",
+                quantized_mistral_experts.weight_scale_inv,
+            )
+        return quantized_mistral_experts
+
+    def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
+        """Forward pass of the block-scale quantized MixtralExperts.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        indptr : nn.Tensor
+            The indptr tensor of group gemm, with shape of [num_experts + 1,].
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor.
+        """
+        if indptr.ndim == 2:
+            # The input is for single token, which does not need group gemm
+            # and can be specialized.
+            expert_indices = indptr
+            assert expert_indices.shape[0] == 1
+            return moe_matmul.dequantize_block_scale_float8_gemv(
+                x,
+                self.weight,
+                self.weight_scale_inv,
+                expert_indices,
+                self.block_size,
+                x.dtype,
+            )
+
+        x_fp8, x_scale = rowwise_group_quant_fp8(
+            x, self.block_size[1], self.weight_dtype, transpose_scale=False
+        )
+        if (
+            extern.get_store().cutlass_gemm
+            and tvm.get_global_func(
+                "cutlass.groupwise_scaled_group_gemm_e4m3fn_e4m3fn", allow_missing=True
+            )
+            is not None
+        ):
+            x = cutlass.fp8_groupwise_scaled_group_gemm(
+                x_fp8,
+                x_scale,
+                self.weight,
+                self.weight_scale_inv,
+                indptr,
+                self.block_size,
+                x.dtype,
+            )
+        else:
+            x = triton.fp8_groupwise_scaled_group_gemm(
+                x_fp8,
+                x_scale,
+                self.weight,
+                self.weight_scale_inv,
+                indptr,
+                self.block_size,
+                x.dtype,
+            )
+        return x
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype
+
+
+def rowwise_group_quant_fp8(
+    x: nn.Tensor,
+    group_size: int,
+    dtype: Literal["float8_e4m3fn", "float8_e5m2"],
+    transpose_scale: bool,
+    eps: float = 1e-10,
+    keep_first_batch_dim: bool = False,
+) -> Tuple[nn.Tensor, nn.Tensor]:  # noqa: UP006
+    """Rowwise group quantization of fp8 tensor.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor.
+
+    group_size : int
+        The group size per row for quantization.
+
+    transpose_scale : bool
+        Whether return the transposed scales or not.
+
+    Returns
+    -------
+    x_fp8 : nn.Tensor
+        The quantized tensor.
+
+    x_scale : nn.Tensor
+        The scales of the quantized tensor.
+        If transpose_scale is True, the shape is
+        (*x.shape[:-2], ceildiv(x.shape[-1], group_size), x.shape[-2]).
+        Otherwise, the shape is (*x.shape[:-1], ceildiv(x.shape[-1], group_size)).
+    """
+    assert x.ndim >= 2
+    assert group_size > 0
+
+    def quantize(x: te.Tensor):
+        num_group = tirx.ceildiv(x.shape[-1], group_size)
+        max_abs_shape = (*x.shape[:-1], num_group)
+        max_abs_reduce_axis = te.reduce_axis((0, group_size), name="r")
+        scale_dtype = "float32"
+        max_abs = te.compute(
+            shape=max_abs_shape,
+            fcompute=lambda *idx: te.max(
+                tirx.if_then_else(
+                    idx[-1] * group_size + max_abs_reduce_axis < x.shape[-1],
+                    tirx.Max(
+                        te.abs(
+                            x(*idx[:-1], idx[-1] * group_size + max_abs_reduce_axis).astype(
+                                scale_dtype
+                            )
+                        ),
+                        eps,
+                    ),
+                    tirx.min_value(scale_dtype),
+                ),
+                axis=max_abs_reduce_axis,
+            ),
+            name="max_abs",
+        )
+        assert dtype in ["float8_e4m3fn", "float8_e5m2"]
+        fp8_max = 448.0 if dtype == "float8_e4m3fn" else 57344.0
+        fp8_min = -fp8_max
+        scale = te.compute(
+            shape=max_abs_shape,
+            fcompute=lambda *idx: max_abs(*idx) / tirx.const(fp8_max, scale_dtype),
+            name="scale",
+        )
+        x_quantized = te.compute(
+            shape=x.shape,
+            fcompute=lambda *idx: tirx.max(
+                tirx.min(
+                    x(*idx).astype(scale_dtype) / scale(*idx[:-1], idx[-1] // group_size),
+                    fp8_max,
+                ),
+                fp8_min,
+            ).astype(dtype),
+            name="x_quantized",
+        )
+        if transpose_scale:
+            if not keep_first_batch_dim:
+                scale = te.compute(
+                    shape=(num_group, *x.shape[:-1]),
+                    fcompute=lambda *idx: scale(*idx[1:], idx[0]),
+                    name="scale",
+                )
+            else:
+                assert len(x.shape) > 2
+                scale = te.compute(
+                    shape=(x.shape[0], num_group, *x.shape[1:-1]),
+                    fcompute=lambda *idx: scale(idx[0], *idx[2:], idx[1]),
+                    name="scale",
+                )
+        return x_quantized, scale
+
+    x_quantized, scale = nn.tensor_expr_op(quantize, name_hint="rowwise_group_quant_fp8", args=[x])
+    return x_quantized, scale
+
+
+def static_activation_group_quant_fp8(
+    x: nn.Tensor,
+    activation_scale: nn.Tensor,
+    group_size: int,
+    dtype: Literal["float8_e4m3fn", "float8_e5m2"],
+) -> nn.Tensor:
+    """Quantize activations with a pre-computed scale."""
+
+    assert activation_scale.ndim == 1
+
+    def quantize(x: te.Tensor, scale: te.Tensor):
+        fp8_max = 448.0 if dtype == "float8_e4m3fn" else 57344.0
+        fp8_min = -fp8_max
+
+        def fcompute(*idx):
+            group_idx = tirx.indexdiv(idx[-1], group_size)
+            return tirx.max(
+                tirx.min(
+                    x(*idx).astype("float32") / scale(group_idx),
+                    fp8_max,
+                ),
+                fp8_min,
+            ).astype(dtype)
+
+        return te.compute(shape=x.shape, fcompute=fcompute, name="static_activation_group_fp8")
+
+    quantized = nn.tensor_expr_op(
+        quantize,
+        name_hint="static_activation_group_fp8",
+        args=[x, activation_scale],
+    )
+    return quantized
+
+
+def broadcast_activation_scale(
+    x: nn.Tensor,
+    activation_scale: nn.Tensor,
+    transpose: bool,
+) -> nn.Tensor:
+    """Broadcast stored activation scales."""
+
+    reshape_shape = (1,) * (x.ndim - 1) + (activation_scale.shape[0],)
+    scale = nn.op.reshape(activation_scale, reshape_shape)
+    scale = nn.op.broadcast_to(scale, (*x.shape[:-1], activation_scale.shape[0]))
+    if transpose:
+        axes = list(range(scale.ndim))
+        axes[-1], axes[-2] = axes[-2], axes[-1]
+        scale = nn.op.permute_dims(scale, axes=axes)
+    return scale
+
+
+def dequantize_float8_groupwise_scaled_gemv(
+    x: nn.Tensor,
+    w: nn.Tensor,
+    w_scale: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+) -> nn.Tensor:
+    """GEMV for FP8 groupwise scaled quantization.
+
+    Parameters
+    ----------
+    x : Tensor
+        The input tensor of shape (k,)
+
+    w : Tensor
+        The quantized weight tensor of shape (n, k)
+
+    w_scale : Tensor
+        The scale tensor of shape
+        (n // block_size[0], k // block_size[1])
+
+    block_size : Tuple[int, int]
+        The block size of the weight tensor.
+
+    out_dtype : str
+        The output dtype of the GEMV computation.
+    """
+    assert x.ndim == 2
+    assert w.ndim == 2
+    assert w_scale.ndim == 2
+    assert x.shape[0] == 1
+    assert x.shape[1] == w.shape[1]
+    _, k = x.shape
+    n, _ = w.shape
+    model_dtype = x.dtype
+    quantize_dtype = w.dtype
+
+    assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0]
+    assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[1]
+
+    def _dequantize(w, s, i, j):
+        return w[i, j].astype(model_dtype) * s[i // block_size[0], j // block_size[1]].astype(
+            model_dtype
+        )
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        x: T.Buffer((1, k), model_dtype),
+        w: T.Buffer((n, k), quantize_dtype),
+        w_scale: T.Buffer(
+            (
+                (n + block_size[0] - 1) // block_size[0],
+                (k + block_size[1] - 1) // block_size[1],
+            ),
+            "float32",
+        ),
+        o: T.Buffer((n,), out_dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        y = T.sblock_alloc_buffer((n, k), model_dtype)
+        for i1, i2 in T.grid(n, k):
+            with T.sblock("dequantize"):
+                i, j = T.axis.remap("SS", [i1, i2])
+                y[i, j] = _dequantize(w, w_scale, i, j)
+        for i1, i2 in T.grid(n, k):
+            with T.sblock("gemv"):
+                i, j = T.axis.remap("SR", [i1, i2])
+                with T.init():
+                    o[i] = T.cast(T.float16(0), out_dtype)
+                o[i] += (x[0, j] * y[i, j]).astype(out_dtype)
+
+    return nn.op.tensor_ir_op(
+        _func,
+        "moe_dequantize_gemv",
+        args=[x, w, w_scale],
+        out=nn.Tensor.placeholder([n], out_dtype),
+    )

python/mlc_llm/quantization/fp8_quantization.py

@@ -0,0 +1,122 @@
+"""Quantization techniques for FP8"""
+
+import numpy as np
+from tvm import relax, runtime
+from tvm.relax.frontend import nn
+
+from mlc_llm.nn import MixtralExperts
+
+from ..op import cutlass, extern, moe_matmul
+from . import per_tensor_quantization as ptq
+from .utils import apply_sharding
+
+
+class FP8PerTensorQuantizeMixtralExperts(ptq.PerTensorQuantizeMixtralExperts):
+    """MixtralExperts with per-tensor quantization in FP8."""
+
+    def __init__(
+        self,
+        num_local_experts,
+        in_features,
+        out_features,
+        config: ptq.PerTensorQuantize,
+        name: str,
+        tensor_parallel_shards=1,
+    ):
+        super().__init__(num_local_experts, in_features, out_features, config, name)
+        self.tensor_parallel_shards = tensor_parallel_shards
+
+    @staticmethod
+    def from_mixtral_experts(
+        src: "MixtralExperts",
+        config: ptq.PerTensorQuantize,
+        name: str,
+    ) -> "FP8PerTensorQuantizeMixtralExperts":
+        """
+        Converts a non-quantized MixtralExperts to a per-tensor quantized MixtralExperts.
+
+        Parameters
+        ----------
+        src : MixtralExperts
+            The non-quantized MixtralExperts
+
+        config : PerTensorQuantize
+            The FP8 quantization weight_config.
+
+        name : str
+            The name of the layer.
+
+        Returns
+        -------
+        ret : MixtralExpertsFP8
+            The per-tensor quantized MixtralExperts.
+        """
+        quantized_mistral_experts = FP8PerTensorQuantizeMixtralExperts(
+            num_local_experts=src.num_local_experts,
+            in_features=src.in_features,
+            out_features=src.out_features,
+            config=config,
+            name=name,
+            tensor_parallel_shards=src.tensor_parallel_shards,
+        )
+
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, f"{shard.name}_q_weight", quantized_mistral_experts.q_weight)
+            # scale doesn't need to be sharded since it's the same for all shards
+
+        return quantized_mistral_experts
+
+    def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
+        w = self.q_weight
+
+        if self.config.calibration_mode == "max":
+            _, x_scale = self.config.quantize_float8(
+                x,
+                quantize_dtype=self.config.activation_dtype,
+                storage_dtype=self.config.activation_dtype,
+            )
+            if self.config.tensor_parallel_shards > 1:
+                x_scale = nn.ccl_allreduce(x_scale, "max")
+            x_scale = nn.extern(
+                "mlc_llm.calibration_observer",
+                [f"{self.name}.q_calibration_scale", "max", x_scale],
+                out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype),
+            )
+            x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype)
+            x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype)
+
+        if indptr.ndim == 2:
+            assert indptr.shape[0] == 1
+            return moe_matmul.dequantize_float8_gemv(
+                x, w, self.q_scale, indptr, self.config.weight_dtype
+            )
+
+        if extern.get_store().cutlass_group_gemm:
+            if self.config.calibration_mode == "inference":
+                if self.q_calibration_scale is not None:
+                    x /= self.q_calibration_scale.astype(x.dtype)
+                x_q = nn.op.astype(x, dtype=self.config.activation_dtype)
+                x_scale = self.q_calibration_scale
+
+            scale = (
+                x_scale * self.q_scale
+                if self.q_scale is not None
+                else nn.wrap_nested(
+                    relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))),
+                    "scale",
+                )
+            )
+            return cutlass.group_gemm(
+                x_q, w, indptr, scale, self.config.weight_dtype, self.config.model_dtype
+            )
+        # Note: convert_weight is target agnostic, so a fallback must be provided
+        w = nn.tensor_expr_op(
+            self.config.dequantize_float8,
+            "dequantize",
+            args=[w, self.q_scale, self.config.weight_dtype],
+        )
+        return moe_matmul.group_gemm(x, w, indptr)
+
+
+ptq.PerTensorQuantizeMixtralExperts._IMPL["fp8"] = FP8PerTensorQuantizeMixtralExperts

python/mlc_llm/quantization/ft_quantization.py

@@ -0,0 +1,404 @@
+"""The FasterTransformer quantization config"""
+
+from dataclasses import dataclass
+from typing import Any, Callable, List, Literal, Optional, Tuple  # noqa: UP035
+
+import tvm
+from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx
+from tvm.relax.frontend import nn
+from tvm.runtime import Tensor
+from tvm.s_tir import dlight as dl
+from tvm.target import Target
+
+from ..loader import QuantizeMapping
+from ..op import faster_transformer_dequantize_gemm
+from ..support import logging
+from ..support.auto_target import detect_cuda_arch_list
+from ..support.style import bold
+from .group_quantization import (
+    GroupQuantize,
+    GroupQuantizeEmbedding,
+    GroupQuantizeLinear,
+)
+from .utils import is_final_fc, is_moe_gate
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FTQuantize:
+    """Configuration for FasterTransformer quantization"""
+
+    name: str
+    kind: str
+    quantize_dtype: Literal["int4", "int8"]
+    storage_dtype: Literal["int8"]
+    model_dtype: Literal["float16"]
+    group_size: Optional[int] = None
+
+    num_elem_per_storage: int = 0
+    max_int_value: int = 0
+
+    def fallback_group_quantize(self) -> GroupQuantize:
+        """
+        The fallback group quantization config for other parameters.
+
+        Returns
+        ------
+        quantize: GroupQuantize
+            The group quantization config to fallback.
+        """
+        return GroupQuantize(
+            name=self.name,
+            kind="group-quant",
+            group_size=32,  # hardcoded to 32 as only supporting int4 quantization
+            quantize_dtype=self.quantize_dtype,
+            storage_dtype="uint32",
+            model_dtype=self.model_dtype,
+            linear_weight_layout="NK",
+        )
+
+    def __post_init__(self):
+        assert self.kind == "ft-quant"
+        quantize_dtype = DataType(self.quantize_dtype)
+        storage_dtype = DataType(self.storage_dtype)
+        assert self.quantize_dtype in ["int4", "int8"]
+        assert storage_dtype.type_code == DataTypeCode.INT
+        assert self.model_dtype == "float16"
+        assert self.group_size in [None, 64, 128]
+        if storage_dtype.bits < quantize_dtype.bits:
+            raise ValueError("Storage unit should be greater or equal to quantized element")
+
+        self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits
+        self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1
+        self._quantize_func_cache = {}
+
+    def quantize_model(
+        self,
+        model: nn.Module,
+        quant_map: QuantizeMapping,
+        name_prefix: str,
+    ) -> nn.Module:
+        """
+        Quantize model with FasterTransformer quantization
+
+        Parameters
+        ----------
+        model : nn.Module
+            The non-quantized nn.Module.
+
+        quant_map : QuantizeMapping
+            The quantize mapping with name mapping and func mapping.
+
+        name_prefix : str
+            The name prefix for visited weight.
+
+        Returns
+        -------
+        ret : nn.Module
+            The quantized nn.Module.
+        """
+
+        class _Mutator(nn.Mutator):
+            def __init__(self, config: FTQuantize, quant_map: QuantizeMapping) -> None:
+                super().__init__()
+                self.config = config
+                self.quant_map = quant_map
+
+            def visit_module(self, name: str, node: nn.Module) -> Any:
+                """
+                The visiting method for FasterTransformer quantization of nn.Module nodes.
+
+                Parameters
+                ----------
+                name : str
+                    The name of the current node.
+
+                node : nn.Module
+                    The current node of nn.Module to mutate.
+
+                Returns
+                ------
+                ret_node: Any
+                    The new node to replace current node.
+                """
+                if isinstance(node, nn.Linear):
+                    weight_name = f"{name}.weight"
+                    self.quant_map.param_map[weight_name] = [
+                        f"{name}.q_weight",
+                        f"{name}.q_scale",
+                    ]
+                    if (
+                        is_final_fc(name)
+                        or node.out_dtype == "float32"
+                        or (self.config.quantize_dtype == "int4" and node.out_features % 8 != 0)
+                        or (self.config.quantize_dtype == "int8" and node.out_features % 4 != 0)
+                    ):
+                        # Under any of the conditions we fall back to GroupQuantize
+                        # For `is_final_fc()` see https://github.com/mlc-ai/mlc-llm/issues/1723
+                        # If simply skipping lm_head quantization degrades performance
+                        # Other requirements are from CUTLASS
+                        logger.info(
+                            'Fallback to GroupQuantize for nn.Linear: "%s", '
+                            + "weight.shape: %s, out_dtype: %s",
+                            bold(name),
+                            node.weight.shape,
+                            node.out_dtype,
+                        )
+                        group_quantize = self.config.fallback_group_quantize()
+                        self.quant_map.map_func[weight_name] = group_quantize.quantize_weight
+                        return GroupQuantizeLinear.from_linear(node, group_quantize)
+                    if not is_moe_gate(name, node):
+                        self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                        return FTQuantizeLinear.from_linear(node, self.config)
+                if isinstance(node, nn.Embedding):
+                    weight_name = f"{name}.weight"
+                    self.quant_map.param_map[weight_name] = [
+                        f"{name}.q_weight",
+                        f"{name}.q_scale",
+                    ]
+                    group_quantize = self.config.fallback_group_quantize()
+                    self.quant_map.map_func[weight_name] = group_quantize.quantize_weight
+                    return GroupQuantizeEmbedding.from_embedding(node, group_quantize)
+                return self.visit(name, node)
+
+        model.to(dtype=self.model_dtype)
+        mutator = _Mutator(self, quant_map)
+        model = mutator.visit(name_prefix, model)
+        return model
+
+    def quantize_weight(self, weight: Tensor) -> List[Tensor]:  # noqa: UP006
+        """
+        Quantize weight with FasterTransformer quantization
+
+        Parameters
+        ----------
+        weight : Tensor
+            The original weight.
+
+        Returns
+        ------
+        ret: List[Tensor]
+            The list of FasterTransformer quantized weights.
+        """
+        assert tvm.get_global_func("relax.ext.cutlass", True), (
+            "Cutlass should be enabled in TVM runtime to quantize weight, "
+            "but not enabled in current TVM runtime environment. "
+            "To enable Cutlass in TVM runtime, please `set(USE_CUTLASS ON)` "
+            "in config.cmake when compiling TVM from source"
+        )
+        assert len(weight.shape) == 2
+        device = weight.device
+        device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+        if device_type == "cuda":
+            target = Target.current()
+            if target is None:
+                target = Target.from_device(device)
+            with target:
+
+                def _create_quantize_func() -> IRModule:
+                    bb = relax.BlockBuilder()
+                    weight_var = relax.Var("weight", relax.TensorType(weight.shape, weight.dtype))
+                    with bb.function(name="main", params=[weight_var]):
+                        with bb.dataflow():
+                            lv0 = bb.emit_te(self._quantize, weight_var)
+                            lv1 = bb.normalize(lv0[0])
+                            lv2 = bb.emit(
+                                relax.call_pure_packed(
+                                    "cutlass.ft_preprocess_weight",
+                                    lv1,
+                                    detect_cuda_arch_list(target=target)[0],
+                                    DataType(self.quantize_dtype).bits == 4,
+                                    ty_args=lv1.ty,
+                                )
+                            )
+                            gv = bb.emit_output(relax.Tuple([lv2, lv0[1]]))
+                        bb.emit_func_output(gv)
+                    return bb.finalize()
+
+                def _compile_quantize_func(mod: IRModule) -> Callable:
+                    mod = dl.ApplyDefaultSchedule(
+                        dl.gpu.Reduction(),
+                        dl.gpu.GeneralReduction(),
+                        dl.gpu.Fallback(),
+                    )(mod)
+                    ex = relax.build(mod, target=target)
+                    vm = relax.VirtualMachine(ex, device)
+                    return vm["main"]
+
+                key = str(
+                    (
+                        int(weight.shape[0]),
+                        int(weight.shape[1]),
+                        weight.dtype,
+                        device_type,
+                    )
+                )
+                quantize_func = self._quantize_func_cache.get(key, None)
+                if quantize_func is None:
+                    logger.info("Compiling quantize function for key: %s", key)
+                    quantize_func = _compile_quantize_func(_create_quantize_func())
+                    self._quantize_func_cache[key] = quantize_func
+                data = quantize_func(weight)
+                return data
+        else:
+            raise NotImplementedError(f"Device type {device_type} is not supported")
+
+    def _quantize(
+        self,
+        weight: te.Tensor,
+    ) -> Tuple[te.Tensor, te.Tensor]:  # noqa: UP006
+        """FasterTransformer quantization for weight tensor, defined in tensor expression."""
+        assert len(weight.shape) == 2
+        n, k = weight.shape
+
+        cur_group_size = k if not self.group_size else self.group_size
+        scale_shape = (tirx.ceildiv(k, cur_group_size), n)
+        r = te.reduce_axis((0, cur_group_size), name="r")
+
+        max_abs = te.compute(
+            shape=scale_shape,
+            fcompute=lambda j, i: te.max(
+                tirx.if_then_else(
+                    j * cur_group_size + r < k,
+                    te.abs(weight[i, j * cur_group_size + r]),
+                    te.min_value(self.model_dtype),
+                ),
+                axis=r,
+            ),
+            name="max_abs_value",
+        )
+        max_int = tirx.const(self.max_int_value, self.model_dtype)
+        scale = te.compute(
+            scale_shape,
+            lambda i, j: max_abs[i, j].astype(self.model_dtype) / max_int,
+            name="scale",
+        )
+        # compute scaled weight
+        quantize_dtype = DataType(self.quantize_dtype)
+        bin_mask = tirx.const((1 << quantize_dtype.bits) - 1, self.storage_dtype)
+        scaled_weight = te.compute(
+            shape=weight.shape,
+            fcompute=lambda i, j: (
+                tirx.min(
+                    tirx.max(
+                        tirx.round(weight[i, j] / scale[j // cur_group_size, i]),
+                        -max_int - 1,
+                    ),
+                    max_int,
+                ).astype(self.storage_dtype)
+                & bin_mask
+            ),
+        )
+
+        quantized_weight_shape = (k, tirx.ceildiv(n, self.num_elem_per_storage))
+        r = te.reduce_axis((0, self.num_elem_per_storage), name="r")
+        quantized_weight = te.compute(
+            shape=quantized_weight_shape,
+            fcompute=lambda j, i: tirx.sum(
+                tirx.if_then_else(
+                    i * self.num_elem_per_storage + r < n,
+                    scaled_weight[i * self.num_elem_per_storage + r, j]
+                    << (
+                        r.astype(self.storage_dtype)
+                        * tirx.const(quantize_dtype.bits, self.storage_dtype)
+                    ),
+                    tirx.const(0, self.storage_dtype),
+                ),
+                axis=r,
+            ),
+            name="weight",
+        )
+
+        return quantized_weight, scale
+
+
+class FTQuantizeLinear(nn.Module):
+    """An nn.Linear module with FasterTransformer quantization"""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        config: FTQuantize,
+        bias: bool = True,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype
+        self.config = config
+        cur_group_size = in_features if not config.group_size else config.group_size
+        self.q_weight = nn.Parameter(
+            (in_features, tirx.ceildiv(out_features, config.num_elem_per_storage)),
+            config.storage_dtype,
+        )
+        self.q_scale = nn.Parameter(
+            (tirx.ceildiv(in_features, cur_group_size), out_features), config.model_dtype
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                (out_features,), config.model_dtype if out_dtype is None else out_dtype
+            )
+        else:
+            self.bias = None
+
+    @staticmethod
+    def from_linear(src: nn.Linear, config: FTQuantize) -> "FTQuantizeLinear":
+        """
+        Converts a non-quantized nn.Linear to a FasterTransformer quantized FTQuantizeLinear
+
+        Parameters
+        ----------
+        src : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : FTQuantize
+            The FasterTransformer quantization config.
+
+        Returns
+        -------
+        ret : FTQuantizeLinear
+            The FasterTransformer quantized FTQuantizeLinear layer.
+        """
+        quantized_linear = FTQuantizeLinear(
+            in_features=src.in_features,
+            out_features=src.out_features,
+            config=config,
+            bias=getattr(src, "bias", None) is not None,
+            out_dtype=src.out_dtype,
+        )
+        if quantized_linear.bias is not None:
+            quantized_linear.bias.attrs = src.bias.attrs
+        return quantized_linear
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        """
+        Forward method for FasterTransformer quantized linear layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the FasterTransformer quantized linear layer.
+        """
+        return faster_transformer_dequantize_gemm(
+            x, self.q_weight, self.q_scale, self.bias, group_size=self.config.group_size
+        )
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        self.q_weight.to(dtype=dtype)
+        self.q_scale.to(dtype=dtype)
+        if self.bias is not None and self.out_dtype is None:
+            self.bias.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype

python/mlc_llm/quantization/group_quantization.py

@@ -0,0 +1,660 @@
+"""The group quantization config"""
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx, topi
+from tvm.relax.frontend import nn
+from tvm.runtime import Tensor
+
+from mlc_llm.loader import QuantizeMapping
+from mlc_llm.nn import MixtralExperts
+from mlc_llm.support import logging
+
+from .utils import (
+    apply_sharding,
+    compile_quantize_func,
+    convert_uint_to_float,
+    is_final_fc,
+    is_moe_gate,
+    pack_weight,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class GroupQuantize:
+    """Configuration for group quantization"""
+
+    name: str
+    kind: str
+    group_size: int
+    quantize_dtype: Literal["int3", "int4", "int8"]
+    storage_dtype: Literal["uint32"]
+    model_dtype: Literal["float16", "float32", "bfloat16"]
+    linear_weight_layout: Literal["KN", "NK"]
+    quantize_embedding: bool = True
+    quantize_final_fc: bool = True
+
+    num_elem_per_storage: int = 0
+    num_storage_per_group: int = 0
+    max_int_value: int = 0
+    tensor_parallel_shards: int = 0
+
+    def __post_init__(self):
+        assert self.kind == "group-quant"
+        quantize_dtype = DataType(self.quantize_dtype)
+        storage_dtype = DataType(self.storage_dtype)
+        model_dtype = DataType(self.model_dtype)
+        assert quantize_dtype.type_code == DataTypeCode.INT
+        assert storage_dtype.type_code == DataTypeCode.UINT
+        assert model_dtype.type_code in (DataTypeCode.FLOAT, DataTypeCode.BFLOAT)
+        if storage_dtype.bits < quantize_dtype.bits:
+            raise ValueError("Storage unit should be greater or equal to quantized element")
+
+        self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits
+        if self.group_size % self.num_elem_per_storage != 0:
+            raise ValueError("Group size should be divisible by numbers of elements per storage")
+        self.num_storage_per_group = self.group_size // self.num_elem_per_storage
+        self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1
+        self.linear_quant_axis = 0 if self.linear_weight_layout == "KN" else 1
+        self._quantize_func_cache = {}
+
+    def quantize_model(
+        self,
+        model: nn.Module,
+        quant_map: QuantizeMapping,
+        name_prefix: str,
+    ) -> nn.Module:
+        """
+        Quantize model with group quantization
+
+        Parameters
+        ----------
+        model : nn.Module
+            The non-quantized nn.Module.
+
+        quant_map : QuantizeMapping
+            The quantize mapping with name mapping and func mapping.
+
+        name_prefix : str
+            The name prefix for visited weight.
+
+        Returns
+        -------
+        ret : nn.Module
+            The quantized nn.Module.
+        """
+
+        class _Mutator(nn.Mutator):
+            def __init__(self, config: GroupQuantize, quant_map: QuantizeMapping) -> None:
+                super().__init__()
+                self.config = config
+                self.quant_map = quant_map
+
+            def visit_module(self, name: str, node: nn.Module) -> Any:
+                """
+                The visiting method for group quantization of nn.Module nodes.
+
+                Parameters
+                ----------
+                name : str
+                    The name of the current node.
+
+                node : nn.Module
+                    The current node of nn.Module to mutate.
+
+                Returns
+                ------
+                ret_node: Any
+                    The new node to replace current node.
+                """
+                if getattr(node, "no_quantization", False):
+                    return node
+
+                if (
+                    isinstance(node, nn.Linear)
+                    and (not is_final_fc(name) or self.config.quantize_final_fc)
+                    and not is_moe_gate(name, node)
+                ):
+                    weight_name = f"{name}.weight"
+                    self.quant_map.param_map[weight_name] = [
+                        f"{name}.q_weight",
+                        f"{name}.q_scale",
+                    ]
+                    self.quant_map.map_func[weight_name] = partial(
+                        self.config.quantize_weight,
+                        output_transpose=self.config.linear_weight_layout == "KN",
+                    )
+                    return GroupQuantizeLinear.from_linear(node, self.config)
+                if isinstance(node, nn.Embedding) and self.config.quantize_embedding:
+                    weight_name = f"{name}.weight"
+                    self.quant_map.param_map[weight_name] = [
+                        f"{name}.q_weight",
+                        f"{name}.q_scale",
+                    ]
+                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                    return GroupQuantizeEmbedding.from_embedding(node, self.config)
+                if isinstance(node, MixtralExperts):
+                    weight_name = f"{name}.weight"
+                    self.quant_map.param_map[weight_name] = [
+                        f"{name}.q_weight",
+                        f"{name}.q_scale",
+                    ]
+                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                    return GroupQuantizeMixtralExperts.from_mixtral_experts(node, self.config)
+                return self.visit(name, node)
+
+        model.to(dtype=self.model_dtype)
+        mutator = _Mutator(self, quant_map)
+        model = mutator.visit(name_prefix, model)
+        return model
+
+    def _dequantize(
+        self,
+        weight: te.Tensor,
+        scale: te.Tensor,
+        axis: int,
+        out_shape: Optional[List[tirx.Expr]] = None,  # noqa: UP006
+    ):
+        tir_max_int = tirx.const(self.max_int_value, self.model_dtype)
+        float_weight = convert_uint_to_float(
+            weight,
+            DataType(self.quantize_dtype).bits,
+            self.num_elem_per_storage,
+            self.storage_dtype,
+            self.model_dtype,
+            axis=axis,
+            out_shape=out_shape,
+        )
+        if out_shape is None:
+            out_shape = weight.shape
+            out_shape[axis] *= self.num_elem_per_storage
+        axis = axis if axis >= 0 else len(out_shape) + axis
+        return te.compute(
+            shape=out_shape,
+            fcompute=lambda *idx: tirx.Mul(
+                tirx.Sub(
+                    float_weight(*idx),
+                    tir_max_int,
+                ),
+                scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :]),
+            ),
+            name="dequantize",
+        )
+
+    def quantize_weight(
+        self, weight: Tensor, axis: int = -1, output_transpose: bool = False
+    ) -> List[Tensor]:  # noqa: UP006
+        """
+        Quantize weight with group quantization
+
+        Parameters
+        ----------
+        weight : Tensor
+            The original weight.
+
+        axis : int
+            The group axis.
+
+        output_transpose : bool
+            Whether to transpose the output quantized weight. Only 2D weight is supported.
+
+        Returns
+        ------
+        ret: List[Tensor]
+            The list of group quantized weights.
+        """
+        device = weight.device
+        device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+        axis = axis if axis >= 0 else len(weight.shape) + axis
+
+        def _create_quantize_func() -> IRModule:
+            bb = relax.BlockBuilder()
+            weight_var = relax.Var("weight", relax.TensorType(weight.shape, weight.dtype))
+            with bb.function(name="main", params=[weight_var]):
+                with bb.dataflow():
+                    lv = bb.emit_te(self._quantize, weight_var, axis, output_transpose)
+                    gv = bb.emit_output(lv)
+                bb.emit_func_output(gv)
+            return bb.finalize()
+
+        key = (
+            f"({weight.shape}, {weight.dtype}, {device_type}, "
+            f"axis={axis}, output_transpose={output_transpose})"
+        )
+        quantize_func = self._quantize_func_cache.get(key, None)
+        if quantize_func is None:
+            logger.info("Compiling quantize function for key: %s", key)
+            quantize_func = compile_quantize_func(_create_quantize_func(), device=device)
+            self._quantize_func_cache[key] = quantize_func
+        return quantize_func(weight)
+
+    def _quantize(
+        self,
+        weight: te.Tensor,
+        axis: int = -1,
+        output_transpose: bool = False,
+    ) -> Tuple[te.Tensor, te.Tensor]:  # noqa: UP006
+        """Group quantization for weight tensor, defined in tensor expression."""
+        max_int = tirx.const(self.max_int_value, self.model_dtype)
+        shape = weight.shape
+        axis = axis if axis >= 0 else len(shape) + axis
+        k = shape[axis]
+        # compute scale per group
+        r = te.reduce_axis((0, self.group_size), name="r")
+        num_group = tirx.ceildiv(k, self.group_size)
+        scale_shape = (*shape[:axis], num_group, *shape[axis + 1 :])
+        max_abs = te.compute(
+            shape=scale_shape,
+            fcompute=lambda *idx: te.max(
+                tirx.if_then_else(
+                    idx[axis] * self.group_size + r < k,
+                    te.abs(
+                        weight(
+                            *idx[:axis],
+                            idx[axis] * self.group_size + r,
+                            *idx[axis + 1 :],
+                        )
+                    ),
+                    te.min_value(self.model_dtype),
+                ),
+                axis=r,
+            ),
+            name="max_abs_value",
+        )
+        scale = te.compute(
+            scale_shape,
+            lambda *idx: max_abs(*idx).astype(self.model_dtype) / max_int,
+            name="scale",
+        )
+        # compute scaled weight
+        scaled_weight = te.compute(
+            shape=weight.shape,
+            fcompute=lambda *idx: tirx.min(
+                tirx.max(
+                    tirx.round(
+                        weight(*idx)
+                        / scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :])
+                        + max_int
+                    ),
+                    tirx.const(0, self.model_dtype),
+                ),
+                max_int * 2,
+            ).astype(self.storage_dtype),
+        )
+        # compute quantized weight per storage
+        num_storage = self.num_storage_per_group * num_group
+        quantized_weight_shape = (*shape[:axis], num_storage, *shape[axis + 1 :])
+        quantized_weight = pack_weight(
+            scaled_weight,
+            axis=axis,
+            num_elem_per_storage=self.num_elem_per_storage,
+            weight_dtype=self.quantize_dtype,
+            storage_dtype=self.storage_dtype,
+            out_shape=quantized_weight_shape,
+        )
+        if output_transpose:
+            if len(quantized_weight.shape) != 2 or len(scale.shape) != 2:
+                raise ValueError(
+                    "Does not support transpose output quantized weight with ndim != 2"
+                )
+            quantized_weight = topi.transpose(quantized_weight)
+            scale = topi.transpose(scale)
+        return quantized_weight, scale
+
+
+class GroupQuantizeLinear(nn.Module):
+    """An nn.Linear module with group quantization"""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: Union[int, tirx.Var],
+        config: GroupQuantize,
+        bias: bool = True,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype
+        self.config = config
+        num_group = tirx.ceildiv(in_features, config.group_size)
+        num_shards = config.tensor_parallel_shards
+        if num_shards > 1 and (in_features * num_shards // config.group_size) % num_shards != 0:
+            raise ValueError(
+                f"The linear dimension {in_features * num_shards} has "
+                f"{in_features * num_shards // config.group_size} groups under group size "
+                f"{config.group_size}. The groups cannot be evenly distributed on "
+                f"{num_shards} GPUs.\n"
+                "Possible solutions: reduce number of GPUs, or use quantization with smaller "
+                "group size."
+            )
+        if config.linear_weight_layout == "KN":
+            self.q_weight = nn.Parameter(
+                (config.num_storage_per_group * num_group, out_features),
+                config.storage_dtype,
+            )
+            self.q_scale = nn.Parameter((num_group, out_features), config.model_dtype)
+        else:
+            self.q_weight = nn.Parameter(
+                (out_features, config.num_storage_per_group * num_group),
+                config.storage_dtype,
+            )
+            self.q_scale = nn.Parameter((out_features, num_group), config.model_dtype)
+        if bias:
+            self.bias = nn.Parameter(
+                (out_features,), config.model_dtype if out_dtype is None else out_dtype
+            )
+        else:
+            self.bias = None
+
+    @staticmethod
+    def from_linear(src: nn.Linear, config: GroupQuantize) -> "GroupQuantizeLinear":
+        """
+        Converts a non-quantized nn.Linear to a group quantized GroupQuantizeLinear
+
+        Parameters
+        ----------
+        src : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : GroupQuantize
+            The group quantization config.
+
+        Returns
+        -------
+        ret : GroupQuantizeLinear
+            The group quantized GroupQuantizeLinear layer.
+        """
+        # For dynamic shape, src.out_features is `"name"`; src.weight.shape[0] is `tirx.Var("name")`
+        out_features, in_features = src.weight.shape
+        quantized_linear = GroupQuantizeLinear(
+            in_features=in_features,
+            out_features=out_features,
+            config=config,
+            bias=getattr(src, "bias", None) is not None,
+            out_dtype=src.out_dtype,
+        )
+        if quantized_linear.bias is not None:
+            quantized_linear.bias.attrs = src.bias.attrs
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight)
+            apply_sharding(shard, f"{shard.name}_q_scale", quantized_linear.q_scale)
+        return quantized_linear
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        """
+        Forward method for group quantized linear layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the group quantized linear layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, scale: self.config._dequantize(
+                weight,
+                scale,
+                axis=self.config.linear_quant_axis,
+                out_shape=(
+                    [
+                        (
+                            tirx.IntImm("int64", self.out_features)
+                            if isinstance(self.out_features, int)
+                            else weight.shape[0]
+                        ),  # Reuse same tirx.Var for symbolic shape (after Exporter)
+                        tirx.IntImm("int64", self.in_features),
+                    ]
+                    if self.config.linear_weight_layout == "NK"
+                    else [
+                        tirx.IntImm("int64", self.in_features),
+                        (
+                            tirx.IntImm("int64", self.out_features)
+                            if isinstance(self.out_features, int)
+                            else weight.shape[1]
+                        ),  # Reuse same tirx.Var for symbolic shape (after Exporter)
+                    ]
+                ),
+            ),
+            name_hint="dequantize",
+            args=[self.q_weight, self.q_scale],
+        )
+        if self.config.linear_weight_layout == "NK":
+            w = nn.op.permute_dims(w)
+        x = nn.op.matmul(x, w, out_dtype=self.out_dtype)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        self.q_weight.to(dtype=dtype)
+        self.q_scale.to(dtype=dtype)
+        if self.bias is not None and self.out_dtype is None:
+            self.bias.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype
+
+
+class GroupQuantizeEmbedding(nn.Module):
+    """An nn.Embedding module with group quantization"""
+
+    def __init__(self, num: Union[int, tirx.Var], dim: int, config: GroupQuantize):
+        self.num = num
+        self.dim = dim
+        self.config = config
+        num_group = tirx.ceildiv(dim, config.group_size)
+        self.q_weight = nn.Parameter(
+            (num, config.num_storage_per_group * num_group), config.storage_dtype
+        )
+        self.q_scale = nn.Parameter((num, num_group), config.model_dtype)
+
+    @staticmethod
+    def from_embedding(embedding: nn.Embedding, config: GroupQuantize) -> "GroupQuantizeEmbedding":
+        """
+        Converts a non-quantized nn.Embedding to a group quantized GroupQuantizeEmbedding
+
+        Parameters
+        ----------
+        linear : nn.Embedding
+            The non-quantized nn.Embedding.
+
+        config : GroupQuantize
+            The group quantization config.
+
+        Returns
+        -------
+        ret : GroupQuantizeEmbedding
+            The group quantized GroupQuantizeEmbedding layer.
+        """
+        num, dim = embedding.weight.shape
+        return GroupQuantizeEmbedding(num, dim, config)
+
+    def forward(self, x: nn.Tensor):
+        """
+        Forward method for group quantized embedding layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the embedding layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, scale: self.config._dequantize(
+                weight,
+                scale,
+                axis=-1,
+                out_shape=[
+                    (
+                        tirx.IntImm("int64", self.num)
+                        if isinstance(self.num, int)
+                        else weight.shape[0]
+                    ),  # Reuse same tirx.Var for symbolic shape (after Exporter)
+                    tirx.IntImm("int64", self.dim),
+                ],
+            ),
+            name_hint="dequantize",
+            args=[self.q_weight, self.q_scale],
+        )
+        if x.ndim == 1:
+            return nn.op.take(w, x, axis=0)
+        return nn.op.reshape(
+            nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0),
+            shape=[*x.shape, self.dim],
+        )
+
+    def lm_head_forward(self, x: nn.Tensor):
+        """The lm_head forwarding, which dequantizes the weight
+        and multiplies it with the input tensor.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the lm_head layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, scale: self.config._dequantize(
+                weight,
+                scale,
+                axis=-1,
+                out_shape=[
+                    (
+                        tirx.IntImm("int64", self.num)
+                        if isinstance(self.num, int)
+                        else weight.shape[0]
+                    ),
+                    tirx.IntImm("int64", self.dim),
+                ],
+            ),
+            name_hint="dequantize",
+            args=[self.q_weight, self.q_scale],
+        )
+        w = nn.op.permute_dims(w)
+        return nn.op.matmul(x, w, out_dtype="float32")
+
+
+class GroupQuantizeMixtralExperts(nn.Module):
+    """An MixtralExperts module with group quantization"""
+
+    def __init__(
+        self,
+        num_local_experts,
+        in_features,
+        out_features,
+        config: GroupQuantize,
+    ):
+        self.num_local_experts = num_local_experts
+        self.in_features = in_features
+        self.out_features = out_features
+        self.config = config
+        num_group = tirx.ceildiv(in_features, config.group_size)
+        self.q_weight = nn.Parameter(
+            (num_local_experts, out_features, config.num_storage_per_group * num_group),
+            config.storage_dtype,
+        )
+        self.q_scale = nn.Parameter(
+            (num_local_experts, out_features, num_group), config.model_dtype
+        )
+        self.quantize_dtype = config.quantize_dtype
+        self.group_size = config.group_size
+        self.dtype = config.model_dtype
+        if config.linear_weight_layout == "KN":
+            raise NotImplementedError("GroupQuantizeMixtralExperts does not support KN layout now.")
+
+    @staticmethod
+    def from_mixtral_experts(
+        src: "MixtralExperts", config: GroupQuantize
+    ) -> "GroupQuantizeMixtralExperts":
+        """
+        Converts a non-quantized MixtralExperts to a group quantized GroupQuantizeMixtralExperts
+
+        Parameters
+        ----------
+        src : MixtralExperts
+            The non-quantized MixtralExperts
+
+        config : GroupQuantize
+            The group quantization config.
+
+        Returns
+        -------
+        ret : GroupQuantizeMixtralExperts
+            The group quantized GroupQuantizeMixtralExperts layer.
+        """
+        quantized_mistral_experts = GroupQuantizeMixtralExperts(
+            num_local_experts=src.num_local_experts,
+            in_features=src.in_features,
+            out_features=src.out_features,
+            config=config,
+        )
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, f"{shard.name}_q_weight", quantized_mistral_experts.q_weight)
+            apply_sharding(shard, f"{shard.name}_q_scale", quantized_mistral_experts.q_scale)
+        return quantized_mistral_experts
+
+    def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
+        """Forward method for group quantized mistral experts.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        indptr: nn.Tensor
+            The indptr tensor
+
+        single_batch_decode: bool
+            Whether to use single-batch decode
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the group quantized mistral experts layer.
+        """
+        from mlc_llm.op import moe_matmul
+
+        assert x.ndim == 2
+        if indptr.ndim == 2:  # single-batch
+            assert indptr.shape[0] == 1
+            return moe_matmul.dequantize_gemv(
+                x,
+                self.q_weight,
+                self.q_scale,
+                indptr,
+                quantize_dtype=self.quantize_dtype,
+                group_size=self.group_size,
+            )
+        assert indptr.ndim == 1
+        return moe_matmul.dequantize_group_gemm(
+            x,
+            self.q_weight,
+            self.q_scale,
+            indptr,
+            quantize_dtype=self.quantize_dtype,
+            indptr_dtype=indptr.dtype,
+            group_size=self.group_size,
+        )

python/mlc_llm/quantization/model_quantization.py

@@ -0,0 +1,157 @@
+"""Quantization factory utilities for model quantization."""
+
+from typing import Any, Callable, Dict, Optional, Tuple, Type  # noqa: UP035
+
+from tvm.relax.frontend import nn
+
+from mlc_llm.loader import QuantizeMapping
+
+from .awq_quantization import AWQQuantize
+from .block_scale_quantization import BlockScaleQuantize
+from .ft_quantization import FTQuantize
+from .group_quantization import GroupQuantize
+from .no_quantization import NoQuantize
+from .per_tensor_quantization import PerTensorQuantize
+from .quantization import Quantization
+
+FuncQuantization = Callable[[Any, Quantization], Tuple[nn.Module, QuantizeMapping]]  # noqa: UP006
+
+
+def make_quantization_functions(
+    model_cls: Type[nn.Module],  # noqa: UP006
+    *,
+    model_ctor: Optional[Callable[[Any], nn.Module]] = None,
+    supports_group_quant: bool = True,
+    supports_ft_quant: bool = True,
+    supports_awq: bool = False,
+    awq_unsupported_message: Optional[str] = None,
+    supports_per_tensor: bool = False,
+    supports_block_scale: bool = False,
+    set_tensor_parallel_shards: bool = True,
+    per_tensor_use_shards: bool = True,
+) -> Dict[str, FuncQuantization]:  # noqa: UP006
+    """Create standard quantization function implementations for a model class."""
+
+    def _create_model(model_config: Any) -> nn.Module:
+        if model_ctor is not None:
+            return model_ctor(model_config)
+        return model_cls(model_config)
+
+    def _no_quant(model_config: Any, quantization: NoQuantize) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        return model, QuantizeMapping({}, {})
+
+    def _group_quant(
+        model_config: Any,
+        quantization: GroupQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        if set_tensor_parallel_shards:
+            if not hasattr(model_config, "tensor_parallel_shards"):
+                raise AttributeError(
+                    "model_config is missing required "
+                    "attribute 'tensor_parallel_shards' for group quantization"
+                )
+            quantization.tensor_parallel_shards = getattr(model_config, "tensor_parallel_shards")
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _ft_quant(model_config: Any, quantization: FTQuantize) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _awq_quant(
+        model_config: Any, quantization: AWQQuantize
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        if awq_unsupported_message is not None:
+            raise NotImplementedError(awq_unsupported_message)
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    def _per_tensor_quant(
+        model_config: Any,
+        quantization: PerTensorQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        kwargs = {}
+        if per_tensor_use_shards:
+            if not hasattr(model_config, "tensor_parallel_shards"):
+                raise AttributeError(
+                    "model_config is missing required attribute "
+                    "'tensor_parallel_shards' for per-tensor quantization"
+                )
+            kwargs["tensor_parallel_shards"] = getattr(model_config, "tensor_parallel_shards")
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+            **kwargs,
+        )
+        return model, quant_map
+
+    def _block_scale_quant(
+        model_config: Any,
+        quantization: BlockScaleQuantize,
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = _create_model(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(model, quant_map, "")
+        return model, quant_map
+
+    quantize_fns: Dict[str, FuncQuantization] = {"no-quant": _no_quant}  # noqa: UP006
+    if supports_group_quant:
+        quantize_fns["group-quant"] = _group_quant
+    if supports_ft_quant:
+        quantize_fns["ft-quant"] = _ft_quant
+    if supports_awq:
+        quantize_fns["awq"] = _awq_quant
+    if supports_per_tensor:
+        quantize_fns["per-tensor-quant"] = _per_tensor_quant
+    if supports_block_scale:
+        quantize_fns["block-scale-quant"] = _block_scale_quant
+    return quantize_fns
+
+
+def make_awq_quant(
+    model_cls: Type[nn.Module],  # noqa: UP006
+) -> Callable[[Any, AWQQuantize], Tuple[nn.Module, QuantizeMapping]]:  # noqa: UP006
+    """Create a standard AWQ quantization function for loaders."""
+
+    def awq_quant(
+        model_config: Any, quantization: AWQQuantize
+    ) -> Tuple[nn.Module, QuantizeMapping]:  # noqa: UP006
+        model = model_cls(model_config)
+        model.to(quantization.model_dtype)
+        quant_map = QuantizeMapping({}, {})
+        model = quantization.quantize_model(
+            model,
+            quant_map,
+            "",
+        )
+        return model, quant_map
+
+    return awq_quant

python/mlc_llm/quantization/no_quantization.py

@@ -0,0 +1,15 @@
+"""The no quantization config"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class NoQuantize:
+    """Configuration for no quantization"""
+
+    name: str
+    kind: str
+    model_dtype: str  # "float16", "float32"
+
+    def __post_init__(self):
+        assert self.kind == "no-quant"

python/mlc_llm/quantization/per_tensor_quantization.py

@@ -0,0 +1,700 @@
+"""The per-tensor quantization config"""
+
+import functools
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, Union  # noqa: UP035
+
+import numpy as np
+from tvm import DataType, DataTypeCode, IRModule, relax, runtime, te, tirx, topi
+from tvm.relax.frontend import nn
+from tvm.runtime import Tensor
+
+from mlc_llm.loader import QuantizeMapping
+from mlc_llm.nn import MixtralExperts
+from mlc_llm.op import cutlass, extern
+from mlc_llm.support import logging
+
+from .utils import (
+    apply_sharding,
+    compile_quantize_func,
+    convert_uint_packed_fp8_to_float,
+    is_final_fc,
+    is_moe_gate,
+    pack_weight,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PerTensorQuantize:
+    """Configuration for per-tensor quantization"""
+
+    name: str
+    kind: str
+    activation_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
+    weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"]
+    storage_dtype: Literal["uint32", "float8_e4m3fn", "float8_e5m2"]
+    model_dtype: Literal["float16"]
+    quantize_embedding: bool = True
+    quantize_final_fc: bool = True
+    quantize_linear: bool = True
+
+    num_elem_per_storage: int = 0
+    max_int_value: int = 0
+    use_scale: bool = True
+    # The calibration mode for quantization. If set to "inference", the model is built for
+    # inference. This should be used after calibration is done.
+    # If set to "max", the model is built for calibration that computes the scale using max value of
+    # the activations.
+    calibration_mode: Literal["inference", "max"] = "inference"
+    tensor_parallel_shards: int = 1
+
+    def __post_init__(self):
+        assert self.kind == "per-tensor-quant"
+        self.num_elem_per_storage = (
+            DataType(self.storage_dtype).bits // DataType(self.weight_dtype).bits
+        )
+        self.max_int_value = int(tirx.max_value(self.weight_dtype).value)
+        self._quantize_func_cache = {}
+
+    def quantize_model(
+        self,
+        model: nn.Module,
+        quant_map: QuantizeMapping,
+        name_prefix: str,
+        tensor_parallel_shards: int,
+    ) -> nn.Module:
+        """
+        Quantize model with per-tensor quantization
+
+        Parameters
+        ----------
+        model : nn.Module
+            The non-quantized nn.Module.
+
+        quant_map : QuantizeMapping
+            The quantize mapping with name mapping and func mapping.
+
+        name_prefix : str
+            The name prefix for visited weight.
+
+        tensor_parallel_shards : int
+            The number of tensor parallel shards.
+
+        Returns
+        -------
+        ret : nn.Module
+            The quantized nn.Module.
+        """
+
+        self.tensor_parallel_shards = tensor_parallel_shards
+
+        class _Mutator(nn.Mutator):
+            def __init__(self, config: PerTensorQuantize, quant_map: QuantizeMapping) -> None:
+                super().__init__()
+                self.config = config
+                self.quant_map = quant_map
+
+            def visit_module(self, name: str, node: nn.Module) -> Any:
+                """
+                The visiting method for per-tensor quantization of nn.Module nodes.
+
+                Parameters
+                ----------
+                name : str
+                    The name of the current node.
+
+                node : nn.Module
+                    The current node of nn.Module to mutate.
+
+                Returns
+                ------
+                ret_node: Any
+                    The new node to replace current node.
+                """
+                weight_name = f"{name}.weight"
+                param_names = (
+                    [f"{name}.q_weight", f"{name}.q_scale"]
+                    if self.config.use_scale
+                    else [
+                        f"{name}.q_weight",
+                    ]
+                )
+                if (
+                    isinstance(node, nn.Linear)
+                    and self.config.quantize_linear
+                    and (not is_final_fc(name) or self.config.quantize_final_fc)
+                    and not is_moe_gate(name, node)
+                ):
+                    self.quant_map.param_map[weight_name] = param_names
+                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                    op = PerTensorQuantizeLinear.from_linear(node, self.config, name)
+                elif isinstance(node, nn.Embedding) and self.config.quantize_embedding:
+                    self.quant_map.param_map[weight_name] = param_names
+                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                    op = PerTensorQuantizeEmbedding.from_embedding(node, self.config)
+                elif isinstance(node, MixtralExperts):
+                    self.quant_map.param_map[weight_name] = param_names
+                    self.quant_map.map_func[weight_name] = self.config.quantize_weight
+                    op = PerTensorQuantizeMixtralExperts.from_mixtral_experts(
+                        node, self.config, name
+                    )
+                else:
+                    return self.visit(name, node)
+
+                if hasattr(op, "q_calibration_scale") and op.q_calibration_scale:
+                    # update quant_map for calibration scale
+                    param_name = f"{name}.q_calibration_scale"
+                    old_map_func = self.quant_map.map_func[weight_name]
+
+                    def map_func(*args, **kwargs):
+                        # placeholder for calibration scale, the actual value will be set after
+                        # calibration.
+                        scale = runtime.empty(
+                            shape=op.q_calibration_scale.shape,
+                            dtype=op.q_calibration_scale.dtype,
+                        )
+                        return [*old_map_func(*args, **kwargs), scale]
+
+                    self.quant_map.param_map[weight_name].append(param_name)
+                    self.quant_map.map_func[weight_name] = map_func
+                return op
+
+        model.to(dtype=self.model_dtype)
+        mutator = _Mutator(self, quant_map)
+        model = mutator.visit(name_prefix, model)
+        return model
+
+    def quantize_weight(self, weight) -> List[Tensor]:  # noqa: UP006
+        """
+        Quantize weight with per-tensor quantization.
+
+        Parameters
+        ----------
+        weight : Tensor
+            The weight to quantize.
+
+        Returns
+        -------
+        ret : List[Tensor]
+            The quantized weight and the scale if use_scale is True.
+        """
+        device = weight.device
+        device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+
+        def _create_quantize_func() -> IRModule:
+            if DataType(self.weight_dtype).type_code in [
+                DataTypeCode.Float8E4M3FN,
+                DataTypeCode.Float8E5M2,
+            ]:
+                quantize_func = functools.partial(
+                    self.quantize_float8,
+                    quantize_dtype=self.weight_dtype,
+                    storage_dtype=self.storage_dtype,
+                )
+            else:
+                assert NotImplementedError()
+
+            class Quantizer(nn.Module):
+                """Quantizer module for per-tensor quantization."""
+
+                def main(self, weight: nn.Tensor):
+                    return quantize_func(weight)
+
+            mod = Quantizer()
+            mod, _ = mod.export_tvm(
+                spec={"main": {"weight": nn.spec.Tensor(weight.shape, weight.dtype)}}
+            )
+            return mod
+
+        key = f"({weight.shape}, {weight.dtype}, {device_type}"
+        quantize_func = self._quantize_func_cache.get(key, None)
+        if quantize_func is None:
+            logger.info("Compiling quantize function for key: %s", key)
+            quantize_func = compile_quantize_func(_create_quantize_func(), device)
+            self._quantize_func_cache[key] = quantize_func
+        return quantize_func(weight)
+
+    def quantize_float8(
+        self,
+        tensor: nn.Tensor,
+        quantize_dtype: str,
+        storage_dtype: str,
+    ) -> Union[Tuple[nn.Tensor], Tuple[nn.Tensor, nn.Tensor]]:  # noqa: UP006
+        """Per-tensor quantization for weight tensor, defined in tensor expression."""
+
+        if self.use_scale:
+            # min_scaling_factor taken from TRT-LLM
+            def _compute_scale(x: te.Tensor) -> te.Tensor:
+                max_abs = topi.max(topi.abs(x))
+                min_scaling_factor = tirx.const(
+                    1.0 / (self.max_int_value * 512.0), self.model_dtype
+                )
+                scale = topi.maximum(
+                    max_abs.astype(self.model_dtype) / self.max_int_value,
+                    min_scaling_factor,
+                ).astype("float32")
+                scale = topi.expand_dims(scale, axis=0)
+                return scale
+
+            scale = nn.tensor_expr_op(_compute_scale, "compute_scale", args=[tensor])
+        else:
+            scale = None
+
+        def _compute_quantized_tensor(weight: te.Tensor, scale: Optional[te.Tensor]) -> te.Tensor:
+            elem_storage_dtype = (
+                f"uint{DataType(quantize_dtype).bits}"
+                if DataType(storage_dtype).type_code == DataTypeCode.UINT
+                else quantize_dtype
+            )
+            scaled_tensor = te.compute(
+                shape=weight.shape,
+                fcompute=lambda *idx: tirx.Cast(
+                    self.storage_dtype,
+                    tirx.reinterpret(
+                        elem_storage_dtype,
+                        tirx.Cast(
+                            quantize_dtype,
+                            weight(*idx) / scale(0) if scale is not None else weight(*idx),
+                        ),
+                    ),
+                ),
+            )
+
+            if quantize_dtype == self.storage_dtype:
+                return scaled_tensor
+
+            packed_weight = pack_weight(
+                scaled_tensor,
+                axis=-1,
+                num_elem_per_storage=self.num_elem_per_storage,
+                weight_dtype=self.weight_dtype,
+                storage_dtype=self.storage_dtype,
+            )
+
+            return packed_weight
+
+        quantized_tensor = nn.tensor_expr_op(
+            _compute_quantized_tensor, "compute_quantized_tensor", args=[tensor, scale]
+        )
+
+        if self.use_scale:
+            return quantized_tensor, scale
+        return (quantized_tensor,)
+
+    def _dequantize(
+        self,
+        q_weight: te.Tensor,
+        scale: Optional[te.Tensor] = None,
+        out_shape: Optional[Sequence[tirx.Expr]] = None,
+    ) -> te.Tensor:
+        if self.use_scale:
+            assert scale is not None
+        if DataType(self.weight_dtype).type_code in [
+            DataTypeCode.Float8E4M3FN,
+            DataTypeCode.Float8E5M2,
+        ]:
+            return self.dequantize_float8(q_weight, scale, self.weight_dtype, out_shape)
+        raise NotImplementedError()
+
+    def dequantize_float8(
+        self,
+        q_tensor: te.Tensor,
+        scale: Optional[te.Tensor],
+        quantize_dtype: str,
+        out_shape: Optional[Sequence[tirx.Expr]] = None,
+    ) -> te.Tensor:
+        """Dequantize a fp8 tensor (input or weight) to higher-precision float."""
+        if quantize_dtype != self.storage_dtype:
+            dequantized_tensor = convert_uint_packed_fp8_to_float(
+                q_tensor,
+                self.num_elem_per_storage,
+                self.storage_dtype,
+                self.model_dtype,
+                quantize_dtype,
+                axis=-1,
+                out_shape=out_shape,
+            )
+        else:
+            dequantized_tensor = q_tensor.astype(self.model_dtype)
+        if scale is not None:
+            dequantized_tensor = dequantized_tensor * scale.astype(dequantized_tensor.dtype)
+        return dequantized_tensor
+
+
+class PerTensorQuantizeLinear(nn.Module):
+    """An nn.Linear module with per-tensor quantization."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: Union[int, tirx.Var],
+        config: PerTensorQuantize,
+        name: str,
+        bias: bool = True,
+        out_dtype: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.out_dtype = out_dtype or config.model_dtype
+        self.config = config
+        self.name = name
+        self.q_weight = nn.Parameter(
+            (out_features, tirx.ceildiv(in_features, config.num_elem_per_storage)),
+            config.storage_dtype,
+        )
+        self.q_calibration_scale = None
+        if config.use_scale:
+            self.q_scale = nn.Parameter((1,), "float32")
+            if config.calibration_mode == "inference":
+                self.q_calibration_scale = nn.Parameter((1,), "float32")
+        else:
+            self.q_scale = None
+        if bias:
+            self.bias = nn.Parameter(
+                (out_features,), config.model_dtype if out_dtype is None else out_dtype
+            )
+        else:
+            self.bias = None
+
+    @classmethod
+    def from_linear(
+        cls, src: nn.Linear, config: PerTensorQuantize, name: str
+    ) -> "PerTensorQuantizeLinear":
+        """
+        Converts a non-quantized nn.Linear to a per-tensor quantized PerTensorQuantizeLinear
+
+        Parameters
+        ----------
+        src : nn.Linear
+            The non-quantized nn.Linear.
+
+        config : PerTensorQuantize
+            The per-tensor quantization config.
+
+        name: str
+            The name of the layer.
+
+        Returns
+        -------
+        ret : PerTensorQuantizeLinear
+            The per-tensor quantized PerTensorQuantizeLinear layer.
+        """
+        out_features, in_features = src.weight.shape
+        quantized_linear = cls(
+            in_features=in_features,
+            out_features=out_features,
+            config=config,
+            name=name,
+            bias=getattr(src, "bias", None) is not None,
+            out_dtype=src.out_dtype,
+        )
+        if quantized_linear.bias is not None:
+            quantized_linear.bias.attrs = src.bias.attrs
+        if "shard_strategy" in src.weight.attrs:
+            shard = src.weight.attrs["shard_strategy"]
+            apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight)
+            # scale doesn't need to be sharded since it's the same for all shards
+        return quantized_linear
+
+    def forward(self, x: nn.Tensor) -> nn.Tensor:
+        """
+        Forward method for per-tensor quantized linear layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the per-tensor quantized linear layer.
+        """
+        # Note: Use calibration scale when calibration is enabled
+        if self.config.calibration_mode == "inference":
+            if self.q_calibration_scale:
+                x /= self.q_calibration_scale.astype(x.dtype)
+            x_q = x.astype(self.config.activation_dtype)
+            x_scale = self.q_calibration_scale
+        elif self.config.calibration_mode == "max":
+            _, x_scale = self.config.quantize_float8(
+                x,
+                quantize_dtype=self.config.activation_dtype,
+                storage_dtype=self.config.storage_dtype,
+            )
+            if self.config.tensor_parallel_shards > 1:
+                x_scale = nn.ccl_allreduce(x_scale, "max")
+            x_scale = nn.extern(
+                "mlc_llm.calibration_observer",
+                [f"{self.name}.q_calibration_scale", "max", x_scale],
+                out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype),
+            )
+            x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype)
+            x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype)
+        else:
+            raise ValueError(f"Unknown calibration mode: {self.config.calibration_mode}")
+
+        if (
+            self.config.weight_dtype == self.config.storage_dtype
+            and self.config.calibration_mode == "inference"
+        ):
+            if (
+                extern.get_store().cutlass_gemm
+                and functools.reduce(lambda x, y: x * y, x_q.shape[:-1]) != 1
+            ):
+                # Dispatch to cutlass kernel for gemm when cutlass is available.
+                scale = (
+                    x_scale * self.q_scale
+                    if self.config.use_scale
+                    else nn.wrap_nested(
+                        relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))),
+                        "scale",
+                    )
+                )
+                return cutlass.fp8_gemm(
+                    x_q,
+                    self.q_weight,
+                    scale,
+                    self.config.weight_dtype,
+                    self.config.model_dtype,
+                )
+            x = nn.op.matmul(x_q, nn.permute_dims(self.q_weight), out_dtype="float32")
+            if self.config.use_scale:
+                scale = x_scale * self.q_scale
+                x = x * scale
+            x = x.astype(self.out_dtype)
+        else:
+            w = nn.op.tensor_expr_op(
+                lambda weight, scale: self.config._dequantize(
+                    weight,
+                    scale,
+                    out_shape=[
+                        (
+                            tirx.IntImm("int64", self.out_features)
+                            if isinstance(self.out_features, int)
+                            else weight.shape[0]
+                        ),
+                        tirx.IntImm("int64", self.in_features),
+                    ],
+                ),
+                "dequantize",
+                args=[self.q_weight, self.q_scale],
+            )
+            x = nn.op.matmul(x, nn.permute_dims(w), out_dtype=self.out_dtype)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def to(self, dtype: Optional[str] = None) -> None:
+        """
+        Override to() such that we do not convert bias if there is an out_dtype.
+        Otherwise, we might run into dtype mismatch when computing x + self.bias.
+        """
+        self.q_weight.to(dtype=dtype)
+        if self.q_scale:
+            self.q_scale.to(dtype=dtype)
+        if self.bias is not None and self.out_dtype is None:
+            self.bias.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype
+
+
+class PerTensorQuantizeEmbedding(nn.Module):
+    """An nn.Embedding module with group quantization"""
+
+    def __init__(self, num: Union[int, tirx.Var], dim: int, config: PerTensorQuantize):
+        self.num = num
+        self.dim = dim
+        self.config = config
+        self.q_weight = nn.Parameter(
+            (num, tirx.ceildiv(dim, config.num_elem_per_storage)), config.storage_dtype
+        )
+        if self.config.use_scale:
+            self.q_scale = nn.Parameter((1,), "float32")
+        else:
+            self.q_scale = None
+
+    @staticmethod
+    def from_embedding(
+        embedding: nn.Embedding, config: PerTensorQuantize
+    ) -> "PerTensorQuantizeEmbedding":
+        """
+        Converts a non-quantized nn.Embedding to a per-tensor quantized PerTensorQuantizeEmbedding
+
+        Parameters
+        ----------
+        linear : nn.Embedding
+            The non-quantized nn.Embedding.
+
+        config : PerTensorQuantize
+            The per-tensor quantization config.
+
+        Returns
+        -------
+        ret : PerTensorQuantizeEmbedding
+            The per-tensor quantized embedding layer.
+        """
+        num, dim = embedding.weight.shape
+        return PerTensorQuantizeEmbedding(num, dim, config)
+
+    def forward(self, x: nn.Tensor):
+        """
+        Forward method for per-tensor quantized embedding layer.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the embedding layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, scale: self.config._dequantize(
+                weight,
+                scale,
+                out_shape=[
+                    (
+                        tirx.IntImm("int64", self.num)
+                        if isinstance(self.num, int)
+                        else weight.shape[0]
+                    ),
+                    tirx.IntImm("int64", self.dim),
+                ],
+            ),
+            "dequantize",
+            args=[self.q_weight, self.q_scale],
+        )
+        if x.ndim == 1:
+            return nn.op.take(w, x, axis=0)
+        return nn.op.reshape(
+            nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0),
+            shape=[*x.shape, self.dim],
+        )
+
+    def lm_head_forward(self, x: nn.Tensor):
+        """The lm_head forwarding, which dequantizes the weight
+        and multiplies it with the input tensor.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the lm_head layer.
+        """
+        w = nn.op.tensor_expr_op(
+            lambda weight, scale: self.config._dequantize(
+                weight,
+                scale,
+                out_shape=[
+                    (
+                        tirx.IntImm("int64", self.num)
+                        if isinstance(self.num, int)
+                        else weight.shape[0]
+                    ),
+                    tirx.IntImm("int64", self.dim),
+                ],
+            ),
+            "dequantize",
+            args=[self.q_weight, self.q_scale],
+        )
+        w = nn.op.permute_dims(w)
+        return nn.op.matmul(x, w, out_dtype="float32")
+
+
+class PerTensorQuantizeMixtralExperts(nn.Module):
+    """An MixtralExperts module with group quantization"""
+
+    _IMPL: ClassVar[Dict[str, Type["PerTensorQuantizeMixtralExperts"]]] = {}  # noqa: UP006
+
+    def __init__(
+        self,
+        num_local_experts,
+        in_features,
+        out_features,
+        config: PerTensorQuantize,
+        name: str,
+    ):
+        self.num_local_experts = num_local_experts
+        self.in_features = in_features
+        self.out_features = out_features
+        self.config = config
+        self.name = name
+        self.q_weight = nn.Parameter(
+            (
+                num_local_experts,
+                out_features,
+                tirx.ceildiv(in_features, config.num_elem_per_storage),
+            ),
+            config.storage_dtype,
+        )
+        self.q_calibration_scale = None
+        if config.use_scale:
+            self.q_scale = nn.Parameter((1,), "float32")
+            if config.calibration_mode == "inference":
+                self.q_calibration_scale = nn.Parameter((1,), "float32")
+        else:
+            self.q_scale = None
+
+    @staticmethod
+    def from_mixtral_experts(
+        src: "MixtralExperts",
+        config: PerTensorQuantize,
+        name: str,
+    ) -> "PerTensorQuantizeMixtralExperts":
+        """
+        Converts a non-quantized MixtralExperts to a per-tensor quantized
+        PerTensorQuantizeMixtralExperts
+
+        Parameters
+        ----------
+        src : MixtralExperts
+            The non-quantized MixtralExperts
+
+        config : PerTensorQuantize
+            The per-tensor quantization config
+
+        name: str
+            The name of the layer.
+
+        Returns
+        -------
+        ret : PerTensorQuantizeMixtralExperts
+            The per-tensor quantized MixtralExperts layer
+        """
+        if DataType(config.weight_dtype).type_code in [
+            DataTypeCode.Float8E4M3FN,
+            DataTypeCode.Float8E5M2,
+        ]:
+            return PerTensorQuantizeMixtralExperts._IMPL["fp8"].from_mixtral_experts(
+                src, config, name
+            )
+        raise NotImplementedError()
+
+    def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor:
+        """Forward method for per-tensor quantized mistral experts.
+
+        Parameters
+        ----------
+        x : nn.Tensor
+            The input tensor.
+
+        indptr: nn.Tensor
+            The indptr tensor
+
+        Returns
+        -------
+        ret : nn.Tensor
+            The output tensor for the per-tensor quantized mistral experts layer.
+        """
+        raise NotImplementedError()

python/mlc_llm/quantization/quantization.py

@@ -0,0 +1,201 @@
+"""A centralized registry of all existing quantization methods and their configurations."""
+
+from typing import Any, Dict  # noqa: UP035
+
+from .awq_quantization import AWQQuantize
+from .block_scale_quantization import BlockScaleQuantize
+from .ft_quantization import FTQuantize
+from .group_quantization import GroupQuantize
+from .no_quantization import NoQuantize
+from .per_tensor_quantization import PerTensorQuantize
+
+Quantization = Any
+"""Quantization is an object that represents an quantization algorithm. It is required to
+have the following fields:
+
+    name : str
+        The name of the quantization algorithm, for example, "q4f16_1".
+
+    kind : str
+        The kind of quantization algorithm, for example, "group-quant", "faster-transformer".
+
+It is also required to have the following method:
+
+    def quantize_model(self, module: nn.Module) -> nn.Module:
+        ...
+
+    def quantize_weight(self, weight: tvm.runtime.Tensor) -> List[tvm.runtime.Tensor]:
+        ...
+"""
+
+QUANTIZATION: Dict[str, Quantization] = {  # noqa: UP006
+    "q0f16": NoQuantize(
+        name="q0f16",
+        kind="no-quant",
+        model_dtype="float16",
+    ),
+    "q0bf16": NoQuantize(
+        name="q0bf16",
+        kind="no-quant",
+        model_dtype="bfloat16",
+    ),
+    "q0f32": NoQuantize(
+        name="q0f32",
+        kind="no-quant",
+        model_dtype="float32",
+    ),
+    "q3f16_0": GroupQuantize(
+        name="q3f16_0",
+        kind="group-quant",
+        group_size=40,
+        quantize_dtype="int3",
+        storage_dtype="uint32",
+        model_dtype="float16",
+        linear_weight_layout="KN",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q3f16_1": GroupQuantize(
+        name="q3f16_1",
+        kind="group-quant",
+        group_size=40,
+        quantize_dtype="int3",
+        storage_dtype="uint32",
+        model_dtype="float16",
+        linear_weight_layout="NK",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4f16_0": GroupQuantize(
+        name="q4f16_0",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="float16",
+        linear_weight_layout="KN",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4f16_1": GroupQuantize(
+        name="q4f16_1",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="float16",
+        linear_weight_layout="NK",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4bf16_0": GroupQuantize(
+        name="q4bf16_0",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="bfloat16",
+        linear_weight_layout="KN",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4bf16_1": GroupQuantize(
+        name="q4bf16_1",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="bfloat16",
+        linear_weight_layout="NK",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4f32_1": GroupQuantize(
+        name="q4f32_1",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="float32",
+        linear_weight_layout="NK",
+        quantize_embedding=True,
+        quantize_final_fc=True,
+    ),
+    "q4f16_2": GroupQuantize(
+        name="q4f16_2",
+        kind="group-quant",
+        group_size=32,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="float16",
+        linear_weight_layout="NK",
+        quantize_embedding=False,
+        quantize_final_fc=False,
+    ),
+    "q4f16_autoawq": AWQQuantize(
+        name="q4f16_autoawq",
+        kind="awq",
+        group_size=128,
+        quantize_dtype="int4",
+        storage_dtype="uint32",
+        model_dtype="float16",
+    ),
+    "q4f16_ft": FTQuantize(
+        name="q4f16_ft",
+        kind="ft-quant",
+        quantize_dtype="int4",
+        storage_dtype="int8",
+        model_dtype="float16",
+    ),
+    "e5m2_e5m2_f16": PerTensorQuantize(
+        name="e5m2_e5m2_f16",
+        kind="per-tensor-quant",
+        activation_dtype="float8_e5m2",
+        weight_dtype="float8_e5m2",
+        storage_dtype="float8_e5m2",
+        model_dtype="float16",
+        quantize_final_fc=False,
+        quantize_embedding=False,
+        quantize_linear=True,
+        use_scale=False,
+    ),
+    "e4m3_e4m3_f16": PerTensorQuantize(
+        name="e4m3_e4m3_f16",
+        kind="per-tensor-quant",
+        activation_dtype="float8_e4m3fn",
+        weight_dtype="float8_e4m3fn",
+        storage_dtype="float8_e4m3fn",
+        model_dtype="float16",
+        quantize_final_fc=False,
+        quantize_embedding=False,
+        quantize_linear=True,
+        use_scale=True,
+        calibration_mode="inference",
+    ),
+    "e4m3_e4m3_f16_max_calibrate": PerTensorQuantize(
+        name="e4m3_e4m3_f16_max_calibrate",
+        kind="per-tensor-quant",
+        activation_dtype="float8_e4m3fn",
+        weight_dtype="float8_e4m3fn",
+        storage_dtype="float8_e4m3fn",
+        model_dtype="float16",
+        quantize_final_fc=False,
+        quantize_embedding=False,
+        quantize_linear=True,
+        use_scale=True,
+        calibration_mode="max",
+    ),
+    "fp8_e4m3fn_bf16_block_scale": BlockScaleQuantize(
+        name="fp8_e4m3fn_bf16_block_scale",
+        kind="block-scale-quant",
+        weight_dtype="float8_e4m3fn",
+        model_dtype="bfloat16",
+    ),
+    "fp8_e4m3fn_bf16_block_scale_static_activation": BlockScaleQuantize(
+        name="fp8_e4m3fn_bf16_block_scale_static_activation",
+        kind="block-scale-quant",
+        weight_dtype="float8_e4m3fn",
+        model_dtype="bfloat16",
+        use_activation_scale=True,
+    ),
+}

python/mlc_llm/quantization/utils.py

@@ -0,0 +1,188 @@
+"""Common utilities for quantization"""
+
+from collections.abc import Sequence
+from typing import Callable, List, Optional  # noqa: UP035
+
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.frontend import nn
+from tvm.runtime import DataType, DataTypeCode
+from tvm.s_tir import dlight as dl
+from tvm.target import Target
+
+from mlc_llm.support import tensor_parallel as tp
+
+
+def convert_uint_to_float(
+    weight: te.Tensor,
+    bits: int,
+    num_elem_per_storage: int,
+    storage_dtype: str,
+    model_dtype: str,
+    axis: int = -1,
+    out_shape: Optional[List[tirx.Expr]] = None,  # noqa: UP006
+    ft_reorder: Optional[bool] = False,
+) -> te.Tensor:
+    """Convert a quantized uint weight to an unquantized float weight."""
+    tir_bin_mask = tirx.const((1 << bits) - 1, storage_dtype)
+    if out_shape is None:
+        out_shape = weight.shape
+        out_shape[axis] *= num_elem_per_storage
+    axis = axis if axis >= 0 else len(out_shape) + axis
+    return te.compute(
+        shape=out_shape,
+        fcompute=lambda *idx: tirx.bitwise_and(
+            tirx.shift_right(
+                weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]),
+                (
+                    (
+                        (idx[axis] % num_elem_per_storage) % 2 * 4
+                        + (idx[axis] % num_elem_per_storage) // 2
+                    )
+                    * bits
+                    if ft_reorder
+                    else (idx[axis] % num_elem_per_storage) * bits
+                ).astype(storage_dtype),
+            ),
+            tir_bin_mask,
+        ).astype(model_dtype),
+    )
+
+
+def is_final_fc(name: str) -> bool:
+    """Determines whether the parameter is the last layer based on its name."""
+    # TODO: use more specious condition to determine final fc
+    return name in ["head", "lm_head", "lm_head.linear", "embed_out"]
+
+
+def is_moe_gate(name: str, node: nn.Linear) -> bool:
+    """Check whether the parameter is the MoE gate layer."""
+    return name.endswith("gate") and isinstance(node.out_features, int) and node.out_features <= 256
+
+
+def compile_quantize_func(mod: IRModule, device) -> Callable:
+    """Compile a quantization function for a given device."""
+    device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+    if device_type in ["cuda", "rocm", "metal", "vulkan", "opencl"]:
+        target = Target.current()
+        if target is None:
+            target = Target.from_device(device)
+        with target:
+            mod = dl.ApplyDefaultSchedule(
+                dl.gpu.Reduction(),
+                dl.gpu.GeneralReduction(),
+                dl.gpu.Fallback(),
+            )(mod)
+    elif device_type == "cpu":
+        target = "llvm"
+        mod = relax.transform.LegalizeOps()(mod)
+    else:
+        raise NotImplementedError(f"Device type {device_type} is not supported")
+    ex = relax.build(mod, target=target)
+    vm = relax.VirtualMachine(ex, device)
+    return vm["main"]
+
+
+def apply_sharding(shard_strategy, name: str, weight: nn.Parameter):
+    """Apply sharding strategy to a weight."""
+    if isinstance(shard_strategy, tp.ShardSingleDim):
+        weight.attrs["shard_strategy"] = tp.ShardSingleDim(
+            name=name,
+            dim=shard_strategy.dim,
+            segs=shard_strategy.segs,
+        )
+    else:
+        raise NotImplementedError(f"Unknowing sharding strategy: {shard_strategy}")
+
+
+def convert_uint_packed_fp8_to_float(
+    weight: te.Tensor,
+    num_elem_per_storage: int,
+    storage_dtype: str,
+    model_dtype: str,
+    quant_dtype: str,
+    axis: int = -1,
+    out_shape: Optional[Sequence[tirx.Expr]] = None,
+) -> te.Tensor:
+    """Unpack a fp8 value from the storage dtype and convert to float."""
+    assert quant_dtype in ["float8_e4m3fn", "float8_e5m2"]
+    assert DataType(storage_dtype).type_code == DataTypeCode.UINT
+    bits = DataType(quant_dtype).bits
+    elem_storage_dtype = DataType(f"uint{bits}")
+    tir_bin_mask = tirx.const((1 << bits) - 1, "uint8")
+    if axis < 0:
+        axis += len(weight.shape)
+    if out_shape is None:
+        out_shape = (
+            *weight.shape[:axis],
+            weight.shape[axis] * num_elem_per_storage,
+            *weight.shape[axis + 1 :],
+        )
+    axis = axis if axis >= 0 else len(out_shape) + axis
+    return te.compute(
+        shape=out_shape,
+        fcompute=lambda *idx: tirx.reinterpret(
+            quant_dtype,
+            tirx.bitwise_and(
+                tirx.shift_right(
+                    weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]),
+                    ((idx[axis] % num_elem_per_storage) * bits).astype(storage_dtype),
+                ).astype(elem_storage_dtype),
+                tir_bin_mask,
+            ),
+        ).astype(model_dtype),
+    )
+
+
+def pack_weight(
+    weight: te.Tensor,
+    axis: int,
+    num_elem_per_storage: int,
+    weight_dtype: str,
+    storage_dtype: str,
+    out_shape: Optional[Sequence[tirx.Expr]] = None,
+):
+    """Convert a tensor to a packed format by packing consecutive bits.
+    This can be useful for sub-byte quantization.
+
+    Parameters
+    ----------
+    weight : te.Tensor
+        The weight
+    axis : int
+        The axis to pack.
+    num_elem_per_storage : int
+        The number of elements per storage.
+    weight_dtype : str
+        The dtype of the input tensor.
+    storage_dtype : str
+        The dtype of the packed tensor.
+    out_shape : Optional[Sequence[tirx.Expr]]
+        The output shape of the packed tensor. Zero-padding is added if needed.
+    """
+    assert weight.dtype == storage_dtype
+    shape = weight.shape
+    if axis < 0:
+        axis += len(shape)
+    k = shape[axis]
+    axis = axis if axis >= 0 else len(shape) + axis
+    if out_shape is None:
+        out_shape = (
+            *shape[:axis],
+            tirx.ceildiv(k, num_elem_per_storage),
+            *shape[axis + 1 :],
+        )
+    r = te.reduce_axis((0, num_elem_per_storage), name="r")
+    packed_weight = te.compute(
+        shape=out_shape,
+        fcompute=lambda *idx: tirx.sum(
+            tirx.if_then_else(
+                idx[axis] * num_elem_per_storage + r < k,
+                weight(*idx[:axis], idx[axis] * num_elem_per_storage + r, *idx[axis + 1 :])
+                << (r * DataType(weight_dtype).bits),
+                tirx.const(0, storage_dtype),
+            ),
+            axis=r,
+        ),
+        name="packed_weight",
+    ).astype(storage_dtype)
+    return packed_weight