Implementation:Mlc ai Mlc llm Quantization Utils

Overview

The Quantization Utils module provides common utility functions shared across all quantization techniques in MLC LLM. Located at python/mlc_llm/quantization/utils.py (189 lines), it implements low-level bit manipulation operations for converting between packed integer and floating-point representations, weight packing, quantization function compilation, sharding strategy application, and layer identification helpers.

Purpose

This module is the shared foundation for the MLC LLM quantization subsystem. It provides:

• Bit-level unpacking of quantized integer weights to floating-point values
• FP8 unpacking from packed storage formats
• Weight packing for sub-byte quantization storage
• Compilation of quantization functions for various hardware targets
• Sharding strategy propagation for tensor-parallel deployments
• Heuristic functions to identify layers that should be excluded from quantization

Key Components

convert_uint_to_float

Converts quantized unsigned integer weights back to floating-point values by extracting individual quantized elements from packed storage:

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[tir.PrimExpr]] = None,
    ft_reorder: Optional[bool] = False,
) -> te.Tensor:
    tir_bin_mask = tir.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: tir.bitwise_and(
            tir.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),
    )

Parameters:

Parameter	Description
weight	The packed weight tensor
bits	Bit width of each quantized element
num_elem_per_storage	Number of quantized elements per storage unit
storage_dtype	Data type of the packed storage (e.g., uint32)
model_dtype	Target floating-point data type (e.g., float16)
axis	Axis along which elements are packed (default: last axis)
out_shape	Optional explicit output shape
ft_reorder	Whether to apply FasterTransformer-style bit reordering

When ft_reorder=True, the bit extraction uses the reordering formula ((idx % N) % 2 * 4 + (idx % N) // 2) * bits instead of the standard sequential (idx % N) * bits.

is_final_fc

Determines whether a layer is the final fully-connected output layer (which should not be quantized):

def is_final_fc(name: str) -> bool:
    return name in ["head", "lm_head", "lm_head.linear", "embed_out"]

is_moe_gate

Checks if a linear layer is a Mixture-of-Experts gate layer (which should not be quantized due to small size):

def is_moe_gate(name: str, node: nn.Linear) -> bool:
    return name.endswith("gate") and isinstance(node.out_features, int) and node.out_features <= 256

compile_quantize_func

Compiles a TVM IR module containing a quantization function into an executable for a given device:

def compile_quantize_func(mod: IRModule, device) -> Callable:
    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"]

For GPU targets (CUDA, ROCm, Metal, Vulkan, OpenCL), it applies DLight scheduling passes for reduction, general reduction, and fallback kernels. For CPU targets, it uses LLVM compilation with operation legalization. The compiled function is returned as a callable from the Relax virtual machine.

apply_sharding

Propagates a sharding strategy to a quantized weight parameter:

def apply_sharding(shard_strategy, name: str, weight: nn.Parameter):
    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}")

Currently only ShardSingleDim strategy is supported, which shards a tensor along a single dimension with specified segment boundaries.

convert_uint_packed_fp8_to_float

Unpacks FP8 values (either float8_e4m3fn or float8_e5m2) from packed unsigned integer storage and converts to float:

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[tir.PrimExpr]] = None,
) -> te.Tensor:
    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 = tir.const((1 << bits) - 1, "uint8")
    # ... axis and shape handling ...
    return te.compute(
        shape=out_shape,
        fcompute=lambda *idx: tir.reinterpret(
            quant_dtype,
            tir.bitwise_and(
                tir.shift_right(
                    weight(...),
                    ((idx[axis] % num_elem_per_storage) * bits).astype(storage_dtype),
                ).astype(elem_storage_dtype),
                tir_bin_mask,
            ),
        ).astype(model_dtype),
    )

Unlike convert_uint_to_float, this function uses tir.reinterpret to correctly interpret the extracted bits as an FP8 value before casting to the model dtype.

pack_weight

Packs a tensor into a compressed format by combining consecutive elements along a specified axis into single storage units:

def pack_weight(
    weight: te.Tensor,
    axis: int,
    num_elem_per_storage: int,
    weight_dtype: str,
    storage_dtype: str,
    out_shape: Optional[Sequence[tir.PrimExpr]] = None,
):
    # ...
    r = te.reduce_axis((0, num_elem_per_storage), name="r")
    packed_weight = te.compute(
        shape=out_shape,
        fcompute=lambda *idx: tir.sum(
            tir.if_then_else(
                idx[axis] * num_elem_per_storage + r < k,
                weight(..., idx[axis] * num_elem_per_storage + r, ...)
                << (r * DataType(weight_dtype).bits),
                tir.const(0, storage_dtype),
            ),
            axis=r,
        ),
        name="packed_weight",
    ).astype(storage_dtype)
    return packed_weight

The function uses a TVM reduction axis to sum bit-shifted elements into packed storage. It handles boundary conditions with zero-padding via tir.if_then_else when the number of elements is not evenly divisible by the packing factor.

Dependencies

• tvm -- TVM compiler framework (IRModule, relax, te, tir)
• tvm.relax.frontend.nn -- Neural network module abstraction
• tvm.runtime.DataType, DataTypeCode -- Data type introspection
• tvm.s_tir.dlight -- DLight scheduling passes for GPU kernel optimization
• tvm.target.Target -- Compilation target specification
• mlc_llm.support.tensor_parallel -- Tensor parallelism sharding strategies

File Location

python/mlc_llm/quantization/utils.py