Implementation:Mlc ai Mlc llm Fuse Dequantize Transpose

Overview

FuseDequantizeTranspose is a TVM compiler pass that fuses dequantization and transpose operations by eliminating the redundant transpose step. In quantized models, weights are stored in a compressed format and must be dequantized before use. When the downstream consumer is a matrix multiplication that expects a transposed weight, the typical pattern involves dequantize followed by transpose. This pass removes the transpose by directly producing dequantized output in the non-transposed layout, since the matmul can then use a non-transposed formulation.

File: python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py

Purpose

The optimization targets the pattern matmul(x, permute_dims(call_tir("dequantize", ...))) where:

• The matmul has a left-hand operand with second-to-last dimension equal to 1 (i.e., a vector-matrix multiply, not a general GeMM)
• The dequantize TIR function internally performs a transpose as its final step

By removing the transpose from the dequantize kernel and adjusting the matmul accordingly, the pass eliminates a full memory read/write of the weight matrix.

Class: FuseDequantizeTranspose

@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose")
class FuseDequantizeTranspose:
    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
        return _DequantizeTransposeFuser(mod).transform()

This is a stateless pass that delegates entirely to the _DequantizeTransposeFuser mutator.

Class: _DequantizeTransposeFuser

@mutator
class _DequantizeTransposeFuser(PyExprMutator):
    def __init__(self, mod: IRModule):
        super().__init__(mod)
        self.mod = mod

    def transform(self) -> IRModule:
        for g_var, func in self.mod.functions_items():
            if isinstance(func, relax.Function):
                updated_func = self.visit_expr(func)
                updated_func = remove_all_unused(updated_func)
                self.builder_.update_func(g_var, updated_func)
        return self.builder_.get()

Pattern Matching in visit_call_

The visit_call_ method checks a chain of conditions to identify the fusion opportunity:

Step 1: Identify matmul with vector LHS

if call.op != tvm.ir.Op.get("relax.matmul"):
    return call
# Do not fuse dequantize-transpose for GeMM
if (
    call.args[0].struct_info.ndim < 2
    or not isinstance(call.args[0].struct_info.shape[-2], tir.IntImm)
    or call.args[0].struct_info.shape[-2].value != 1
):
    return call

The pass only applies when the left operand's second-to-last dimension is statically known to be 1 (vector-matrix multiply). General matrix-matrix multiplications (GeMM) are excluded.

Step 2: Identify transpose (permute_dims) on the RHS

matmul_rhs = self.lookup_binding(call.args[1])
if (
    not isinstance(matmul_rhs, relax.Call)
    or matmul_rhs.op != tvm.ir.Op.get("relax.permute_dims")
    or matmul_rhs.args[0].struct_info.ndim != 2
    or matmul_rhs.attrs.axes is not None
):
    return call

The RHS must be a permute_dims call with default axes (simple transpose) on a 2D tensor.

Step 3: Identify dequantize TIR call as transpose input

transpose_input = self.lookup_binding(matmul_rhs.args[0])
if (
    not isinstance(transpose_input, relax.Call)
    or transpose_input.op != tvm.ir.Op.get("relax.call_tir")
    or not transpose_input.args[0].name_hint.startswith("dequantize")
    or not isinstance(transpose_input.struct_info, relax.TensorStructInfo)
):
    return call

The input to the transpose must be a call_tir whose function name starts with "dequantize".

Step 4: Verify TIR function contains a trailing T_transpose block

dequantize_tir_func = self.mod[transpose_input.args[0]]
if (
    len(dequantize_tir_func.body.block.alloc_buffers) != 1
    or not isinstance(dequantize_tir_func.body.block.body, tir.SeqStmt)
    or len(dequantize_tir_func.body.block.body) != 2
    or not isinstance(dequantize_tir_func.body.block.body[1], tir.For)
    or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tir.SBlockRealize)
    or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose"
):
    return call

The TIR function must have exactly one allocated buffer, a body consisting of exactly two statements (the dequantize computation and a transpose), and the second statement must be a T_transpose block.

Rewrite Logic

When all conditions are satisfied, the pass creates a new TIR function that retains only the dequantization part (the first statement), using the intermediate buffer (before transpose) as the output:

new_func_buffers = [dequantize_tir_func.buffer_map[var] for var in dequantize_tir_func.params]
new_func_buffers[-1] = dequantize_tir_func.body.block.alloc_buffers[0]  # intermediate buffer
new_func = tir.PrimFunc(
    params=new_func_buffers,
    body=tir.SBlockRealize(
        iter_values=[], predicate=True,
        block=tir.SBlock(
            iter_vars=[], reads=[], writes=[], name_hint="root",
            body=dequantize_tir_func.body.block.body[0],  # only dequantize, no transpose
        ),
    ),
)
new_func = tir.stmt_functor.renew_defs(new_func)

The renew_defs call performs a deep copy to avoid IR node sharing between PrimFuncs. The new function is added to the module and the matmul is rewritten to use the dequantized (non-transposed) result directly:

g_var = self.builder_.add_func(new_func, func_name="dequantize")
dequantize_matmul_rhs = self.builder_.emit(
    relax.call_tir(g_var, transpose_input.args[1], out_sinfo=matmul_rhs.struct_info)
)
return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype)

Transformation Summary

Before	After
matmul(x, permute_dims(call_tir("dequantize", inputs)))	matmul(x, call_tir("dequantize_no_transpose", inputs))

The dequantize TIR function is rewritten to output in the transposed layout directly, and the explicit permute_dims is removed.

Limitations

• Only applies to vector-matrix multiplies (LHS second-to-last dim == 1), not general GeMM
• Only handles 2D transpose with default axes
• Requires the dequantize TIR function to have a specific two-statement structure (dequantize + transpose)

Dependencies

• tvm -- Core TVM framework (relax, tir)
• tvm.relax.analysis.remove_all_unused -- Removes unused bindings after rewriting
• tvm.relax.expr_functor -- PyExprMutator and @mutator decorator