chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -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