Repository: mlc-ai/mlc-llm
File: mlc_llm/quantization/ft_quantization.py (Python)

Code before:
                encoding_func(
                    self.nbit,
                    8,
                    group_size=self.group_size,
                    dtype=self.dtype,
                ),
                inputs[0],
                primfunc_name_hint="encode",
            )
            packed_weight = bb.normalize(encoded_data[0])
            encoded_weight = bb.emit(
                relax.call_pure_packed(
                    "cutlass.ft_preprocess_weight",
                    packed_weight,
                    self.sm,
                    self.nbit == 4,
                    sinfo_args=packed_weight.struct_info,
                )
            )
            return bb.emit(relax.Tuple([encoded_weight, encoded_data[1]]))

        return f_quantize

    def get_dequantize_func(
        self,
        param_info: relax.TensorStructInfo,
        qparam_info: List[relax.TensorStructInfo],
    ) -> Optional[FQuantize]:
        return convert_TE_func(
            decoding_func(
                self.nbit,
                storage_nbit=8,
                group_size=self.group_size,
            ),
            func_name="decode",
        )


def encoding_func(nbit: int, storage_nbit: int, group_size: int, dtype: str = "float32"):
    def te_encode_sym(weight: te.Tensor):
        """Encode the weight tensor of shape [N, K] into a quantized weight tensor of shape
        [K, N // float_per_int] and a scale tensor of shape [K // group_size, N]
        """
        n_float_per_int = storage_nbit // nbit
        max_int_value = (1 << (nbit - 1)) - 1

        cur_group_size = weight.shape[1] if group_size == -1 else group_size
        scale_min_shape = (tir.ceildiv(weight.shape[1], cur_group_size), weight.shape[0])
        k = te.reduce_axis((0, cur_group_size), name="k")
        max_abs_value = te.compute(
            shape=scale_min_shape,

Review comment (suggestion):
We need to skip weight preprocessing during quantization, for FT quantization + `num_shards > 1` case.

Code after:
        def f_quantize(bb: relax.BlockBuilder, inputs: List[relax.Expr]):
            encoded_data = bb.emit_te(
                encoding_func(
                    self.nbit,
                    8,
                    group_size=self.group_size,
                    dtype=self.dtype,
                ),
                inputs[0],
                primfunc_name_hint="encode",
            )

            packed_weight = bb.normalize(encoded_data[0])

            if self.do_preprocess:
                encoded_weight = bb.emit(
                    relax.call_pure_packed(
                        "cutlass.ft_preprocess_weight",
                        packed_weight,
                        self.sm,
                        self.nbit == 4,
                        sinfo_args=packed_weight.struct_info,
                    )
                )
            else:
                encoded_weight = packed_weight

            return bb.emit(relax.Tuple([encoded_weight, encoded_data[1]]))

        return f_quantize

    def get_dequantize_func(
        self,
        param_info: relax.TensorStructInfo,
        qparam_info: List[relax.TensorStructInfo],
    ) -> Optional[FQuantize]:
        return convert_TE_func(
            decoding_func(
                self.nbit,
                storage_nbit=8,
                group_size=self.group_size,
            ),
            func_name="decode",
        )


def encoding_func(nbit: int, storage_nbit: int, group_size: int, dtype: str = "float32"):
    def te_encode_sym(weight: te.Tensor):
        """Encode the weight tensor of shape [N, K] into a quantized weight tensor of shape
        [K, N // float_per_int] and a scale tensor of shape [K // group_size, N]
        """