chore: import upstream snapshot with attribution

python/mlc_llm/compiler_pass/__init__.py

@@ -0,0 +1,3 @@
+"""Compiler passes used in MLC LLM."""
+
+from . import pipeline as _pipeline

python/mlc_llm/compiler_pass/attach_cuda_graph_alloc_init_func.py

@@ -0,0 +1,33 @@
+"""The pass that attaches an empty function for initialization."""
+
+import tvm
+from tvm import IRModule, relax
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphAllocInitFunc")
+class AttachCUDAGraphAllocInitFunc:
+    """Attach an empty function for initialization."""
+
+    def __init__(self):
+        pass
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        bb = relax.BlockBuilder(mod)
+        alloc_func_gv = None
+        for gv, _ in mod.functions_items():
+            if gv.name_hint.startswith("cuda_graph_alloc"):
+                assert alloc_func_gv is None
+                alloc_func_gv = gv
+        if alloc_func_gv is None:
+            return mod
+
+        with bb.function("cuda_graph_alloc_init", []):
+            bb.emit_func_output(
+                relax.op.call_builtin_with_ctx(
+                    "vm.builtin.cuda_graph.get_cached_alloc",
+                    args=[alloc_func_gv, relax.prim_value(0)],
+                    ty_args=relax.ObjectType(),
+                )
+            )
+        return bb.finalize()

python/mlc_llm/compiler_pass/attach_embedding_allocator.py

@@ -0,0 +1,39 @@
+"""The pass that attaches embedding allocation function to the IRModule."""
+
+from typing import Any, Dict  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachAllocEmbeddingTensorFunc")
+class AttachAllocEmbeddingTensorFunc:
+    """Attach embedding tensor allocation Relax function to IRModule."""
+
+    def __init__(self, metadata: Dict[str, Any]):  # noqa: UP006
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        embed_func = None
+        for gv, func in mod.functions_items():
+            if gv.name_hint == "embed":
+                embed_func = func
+
+        if embed_func is None:
+            return mod
+
+        hidden_size = embed_func.ret_ty.shape[-1]
+        dtype = relax.DataTypeImm(embed_func.ret_ty.dtype.dtype)
+        bb = relax.BlockBuilder(mod)
+        with bb.function("alloc_embedding_tensor", []):
+            bb.emit_func_output(
+                bb.emit(
+                    relax.op.builtin.alloc_tensor(
+                        relax.ShapeExpr([self.metadata["prefill_chunk_size"], hidden_size]),
+                        dtype,
+                        runtime_device_index=0,
+                    )
+                )
+            )
+        return bb.finalize()

python/mlc_llm/compiler_pass/attach_logit_processor.py

@@ -0,0 +1,285 @@
+"""The pass that attaches logit processor functions to the IRModule."""
+
+import tvm
+from tvm import IRModule
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import (
+    check_thread_limits,
+    get_max_num_threads_per_block,
+)
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachLogitProcessFunc")
+class AttachLogitProcessFunc:
+    """Attach logit processing TIR functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target):
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            The target of the model compilation.
+        """
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        mod = mod.clone()
+        if self.target.kind.name == "llvm":
+            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace_cpu()
+            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace_cpu()
+            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace_cpu()
+        else:
+            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace(self.target)
+            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace(self.target)
+            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace(self.target)
+        return mod
+
+
+def _get_apply_logit_bias_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_logit_bias_inplace(
+        var_logits: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_logit_bias: T.handle,
+    ) -> None:
+        """Function that applies logit bias in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_logit_bias_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        # seq_ids
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32")
+
+        for i in range(num_token):
+            logits[pos2seq_id[i], token_ids[i]] += logit_bias[i]
+
+    return _apply_logit_bias_inplace
+
+
+def _get_apply_logit_bias_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_logit_bias_inplace(
+        var_logits: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_logit_bias: T.handle,
+    ) -> None:
+        """Function that applies logit bias in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_logit_bias_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        # seq_ids
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32")
+
+        for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"):
+            for p1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vp = T.axis.spatial(num_token, p0 * tx + p1)
+                    T.where(p0 * tx + p1 < num_token)
+                    logits[pos2seq_id[vp], token_ids[vp]] += logit_bias[vp]
+
+    return _apply_logit_bias_inplace
+
+
+def _get_apply_penalty_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_penalty_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_token_cnt: T.handle,
+        var_penalties: T.handle,
+    ) -> None:
+        """Function that applies penalties in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_penalty_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32")
+        penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32")
+
+        for token in T.serial(num_token):
+            with T.sblock("block"):
+                vp = T.axis.spatial(num_token, token)
+                logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= (
+                    penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1]
+                )
+                logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else(
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0),
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] * penalties[pos2seq_id[vp], 2],
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] / penalties[pos2seq_id[vp], 2],
+                )
+
+    return _apply_penalty_inplace
+
+
+def _get_apply_penalty_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_penalty_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_token_cnt: T.handle,
+        var_penalties: T.handle,
+    ) -> None:
+        """Function that applies penalties in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_penalty_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32")
+        penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32")
+
+        for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"):
+            for p1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vp = T.axis.spatial(num_token, p0 * tx + p1)
+                    T.where(p0 * tx + p1 < num_token)
+                    # Penalties: (presence_penalty, frequency_penalty, repetition_penalty)
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= (
+                        penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1]
+                    )
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else(
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0),
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]
+                        * penalties[pos2seq_id[vp], 2],
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]
+                        / penalties[pos2seq_id[vp], 2],
+                    )
+
+    return _apply_penalty_inplace
+
+
+def _get_apply_bitmask_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_bitmask_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_bitmask: T.handle,
+    ) -> None:
+        """Function that applies vocabulary masking in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_bitmask_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32")
+
+        for token in T.serial(num_seq * vocab_size):
+            with T.sblock("block"):
+                vs = T.axis.spatial(num_seq, (token) // vocab_size)
+                vv = T.axis.spatial(vocab_size, (token) % vocab_size)
+
+                logits[seq_ids[vs], vv] = T.if_then_else(
+                    (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1,
+                    logits[seq_ids[vs], vv],
+                    T.min_value("float32"),
+                )
+
+    return _apply_bitmask_inplace
+
+
+def _get_apply_bitmask_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_bitmask_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_bitmask: T.handle,
+    ) -> None:
+        """Function that applies vocabulary masking in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_bitmask_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32")
+
+        for fused_s_v_0 in T.thread_binding(0, (num_seq * vocab_size + tx - 1) // tx, "blockIdx.x"):
+            for fused_s_v_1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vs = T.axis.spatial(num_seq, (fused_s_v_0 * tx + fused_s_v_1) // vocab_size)
+                    vv = T.axis.spatial(vocab_size, (fused_s_v_0 * tx + fused_s_v_1) % vocab_size)
+                    T.where(fused_s_v_0 * tx + fused_s_v_1 < num_seq * vocab_size)
+                    logits[seq_ids[vs], vv] = T.if_then_else(
+                        (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1,
+                        logits[seq_ids[vs], vv],
+                        T.min_value("float32"),
+                    )
+
+    return _apply_bitmask_inplace

python/mlc_llm/compiler_pass/attach_sampler.py

@@ -0,0 +1,390 @@
+"""The pass that attaches GPU sampler functions to the IRModule."""
+
+from typing import Dict  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.frontend import nn
+from tvm.script import tirx as T
+
+from mlc_llm.op.batch_spec_verify import batch_spec_verify
+from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachGPUSamplingFunc")
+class AttachGPUSamplingFunc:
+    """Attach GPU sampling functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target, variable_bounds: Dict[str, int]):  # noqa: UP006
+        # Specifically for RWKV workloads, which contains -1 max_seq_len
+        max_batch_size = variable_bounds["batch_size"]
+        self.variable_bounds = {
+            "batch_size": max_batch_size,
+            "num_samples": max_batch_size,
+            "num_positions": 6 * max_batch_size,
+        }
+        self.non_negative_var = ["vocab_size"]
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        target_kind = self.target.kind.name
+        if target_kind not in ["cuda", "vulkan", "metal", "webgpu"]:
+            # Only enable GPU sampling for CUDA, Vulkan, Metal, and WebGPU.
+            return mod
+
+        bb = relax.BlockBuilder(mod)
+        if target_kind == "webgpu":
+            # Only attach functions that do not contain i8s for WebGPU
+            gv_names = [
+                gv.name_hint
+                for gv in [
+                    _attach_argsort_func(bb),
+                    _attach_sample_with_top_p(bb),
+                ]
+            ]
+        else:
+            gv_names = [
+                gv.name_hint
+                for gv in [
+                    _attach_multinomial_sampling_func(bb),
+                    _attach_argsort_func(bb),
+                    _attach_sample_with_top_p(bb),
+                    _attach_take_probs_func(bb),
+                    _attach_batch_verifier(bb),
+                    _attach_renormalize_by_top_p(bb, self.target),
+                ]
+            ]
+
+        mod = bb.finalize()
+        for gv_name in gv_names:
+            mod[gv_name] = (
+                mod[gv_name]
+                .with_attr("tir_var_upper_bound", self.variable_bounds)
+                .with_attr("tir_non_negative_var", self.non_negative_var)
+            )
+        return mod
+
+
+def _attach_multinomial_sampling_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    with bb.function("multinomial_from_uniform", [probs, uniform_samples, sample_indices]):
+        with bb.dataflow():
+            sample_shape = relax.ShapeExpr([num_samples, 1])
+            probs_tensor = nn.wrap_nested(probs, name="probs")
+            uniform_samples_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    uniform_samples,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "float32"),
+                ),
+                name="uniform_samples",
+            )
+            sample_indices_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    sample_indices,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "int32"),
+                ),
+                name="sample_indices",
+            )
+            result_tensor = nn.multinomial_from_uniform(
+                probs_tensor,
+                uniform_samples_tensor,
+                sample_indices_tensor,
+                "int32",
+                name="nn_multinomial_from_uniform",
+            )
+            result = bb.emit(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    result_tensor._expr,
+                    sample_indices.ty.shape,
+                    ty_args=sample_indices.ty,
+                )
+            )
+            output = bb.emit_output(result)
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+def _attach_argsort_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    with bb.function("argsort_probs", [probs]):
+        with bb.dataflow():
+            sorted_indices = bb.emit(relax.op.argsort(probs, descending=True, dtype="int32"))
+            sorted_values = bb.emit_te(
+                lambda unsorted_probs, sorted_indices: te.compute(
+                    (batch_size, vocab_size),
+                    lambda i, j: unsorted_probs[i, sorted_indices[i, j]],
+                    name="take_sorted_probs",
+                ),
+                probs,
+                sorted_indices,
+                primfunc_name_hint="take_sorted_probs",
+            )
+            output = bb.emit_output((sorted_values, sorted_indices))
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+@T.prim_func(s_tir=True)
+def full(var_result: T.handle, value: T.int32):
+    """The filling function for top k."""
+    batch_size = T.int32()
+    result = T.match_buffer(var_result, (batch_size, 1), "int32")
+    for i in T.serial(batch_size):
+        with T.sblock("block"):
+            vi = T.axis.spatial(batch_size, i)
+            result[vi, 0] = value
+
+
+def _attach_sample_with_top_p(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    sorted_probs = relax.Var("sorted_probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    sorted_indices = relax.Var(
+        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
+    )
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
+
+    with bb.function(
+        "sample_with_top_p",
+        [sorted_probs, sorted_indices, uniform_samples, sample_indices, top_p],
+    ):
+        with bb.dataflow():
+            sample_shape = relax.ShapeExpr([num_samples, 1])
+            top_p_shape = relax.ShapeExpr([batch_size, 1])
+            sorted_probs_tensor = nn.wrap_nested(sorted_probs, name="sorted_probs")
+            sorted_indices_tensor = nn.wrap_nested(sorted_indices, name="sorted_indices")
+            uniform_samples_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    uniform_samples,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "float32"),
+                ),
+                name="uniform_samples",
+            )
+            sample_indices_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    sample_indices,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "int32"),
+                ),
+                name="sample_indices",
+            )
+            top_p_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    top_p,
+                    top_p_shape,
+                    ty_args=relax.TensorType(top_p_shape, "float32"),
+                ),
+                name="sample_indices",
+            )
+            top_k_tensor = nn.tensor_ir_op(
+                full,
+                name_hint="full",
+                args=[vocab_size],
+                out=nn.Tensor.placeholder(
+                    [batch_size, 1],
+                    "int32",
+                ),
+            )
+
+            result_tensor = nn.sample_top_p_top_k_from_sorted_prob(
+                sorted_probs_tensor,
+                sorted_indices_tensor,
+                top_p_tensor,
+                top_k_tensor,
+                uniform_samples_tensor,
+                sample_indices_tensor,
+            )
+            result = bb.emit_output(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    result_tensor._expr,
+                    sample_indices.ty.shape,
+                    ty_args=sample_indices.ty,
+                )
+            )
+        gv = bb.emit_func_output(result)
+    return gv
+
+
+def _attach_renormalize_by_top_p(bb: relax.BlockBuilder, target: tvm.target.Target):
+    batch_size = tirx.Var("batch_size", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    num_pivots = 3
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
+    init_pivots = relax.Var("init_pivots", relax.TensorType((batch_size, num_pivots), "float32"))
+    with bb.function("renormalize_by_top_p", [probs, top_p, init_pivots]):
+        with bb.dataflow():
+            cutoff_output = bb.emit(
+                relax.call_tir(
+                    bb.add_func(top_p_pivot(num_pivots, target), "top_p_pivot_cutoff"),
+                    args=[probs, top_p, init_pivots],
+                    out_ty=[top_p.ty, top_p.ty],
+                )
+            )
+            final_pivot = cutoff_output[0]
+            renorm_sum = cutoff_output[1]
+            renormalized_probs = bb.emit_output(
+                relax.call_tir(
+                    bb.add_func(top_p_renorm(target), "top_p_renorm_after_cutoff"),
+                    args=[probs, final_pivot, renorm_sum],
+                    out_ty=probs.ty,
+                )
+            )
+        gv = bb.emit_func_output(renormalized_probs)
+    return gv
+
+
+def _attach_take_probs_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    num_positions = tirx.Var("num_positions", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    unsorted_probs = relax.Var(
+        "unsorted_probs", relax.TensorType((batch_size, vocab_size), "float32")
+    )
+    sorted_indices = relax.Var(
+        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
+    )
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    sampling_results = relax.Var("sampling_result", relax.TensorType((num_samples,), "int32"))
+    top_prob_offsets = relax.Var("lobprob_offsets", relax.TensorType((num_positions,), "int32"))
+
+    @T.prim_func(s_tir=True)
+    def sampler_take_probs_tir(
+        var_unsorted_probs: T.handle,
+        var_sorted_indices: T.handle,
+        var_sample_indices: T.handle,
+        var_sampling_results: T.handle,
+        var_top_prob_offsets: T.handle,
+        var_sampled_values: T.handle,
+        var_top_prob_probs: T.handle,
+        var_top_prob_indices: T.handle,
+    ):
+        batch_size = T.int32()
+        num_samples = T.int32()
+        num_positions = T.int32()
+        vocab_size = T.int32()
+        unsorted_probs = T.match_buffer(var_unsorted_probs, (batch_size, vocab_size), "float32")
+        sorted_indices = T.match_buffer(var_sorted_indices, (batch_size, vocab_size), "int32")
+        sample_indices = T.match_buffer(var_sample_indices, (num_samples,), "int32")
+        sampling_results = T.match_buffer(var_sampling_results, (num_samples,), "int32")
+        top_prob_offsets = T.match_buffer(var_top_prob_offsets, (num_positions,), "int32")
+        sampled_values = T.match_buffer(var_sampled_values, (num_samples,), "float32")
+        top_prob_probs = T.match_buffer(var_top_prob_probs, (num_positions,), "float32")
+        top_prob_indices = T.match_buffer(var_top_prob_indices, (num_positions,), "int32")
+        for i in T.serial(num_positions):
+            with T.sblock("top_prob"):
+                vi = T.axis.spatial(num_positions, i)
+                # Reads are data-dependent gathers; declare full-buffer read
+                # regions explicitly so tirx does not infer data-dependent regions.
+                T.reads(
+                    top_prob_offsets[vi],
+                    sorted_indices[0:batch_size, 0:vocab_size],
+                    unsorted_probs[0:batch_size, 0:vocab_size],
+                )
+                T.writes(top_prob_indices[vi], top_prob_probs[vi])
+                row = T.floordiv(top_prob_offsets[vi], vocab_size)
+                col = T.floormod(top_prob_offsets[vi], vocab_size)
+                top_prob_indices[vi] = sorted_indices[row, col]
+                top_prob_probs[vi] = unsorted_probs[row, sorted_indices[row, col]]
+        for i in T.serial(num_samples):
+            with T.sblock("sample"):
+                vj = T.axis.spatial(num_samples, i)
+                T.reads(
+                    sample_indices[vj],
+                    sampling_results[vj],
+                    unsorted_probs[0:batch_size, 0:vocab_size],
+                )
+                T.writes(sampled_values[vj])
+                sampled_values[vj] = unsorted_probs[sample_indices[vj], sampling_results[vj]]
+
+    args = [
+        unsorted_probs,
+        sorted_indices,
+        sample_indices,
+        sampling_results,
+        top_prob_offsets,
+    ]
+    with bb.function("sampler_take_probs", args):
+        with bb.dataflow():
+            taken_probs_indices = bb.emit_output(
+                relax.call_tir(
+                    bb.add_func(sampler_take_probs_tir, "sampler_take_probs_tir"),
+                    args,
+                    out_ty=[
+                        relax.TensorType((num_samples,), "float32"),
+                        relax.TensorType((num_positions,), "float32"),
+                        relax.TensorType((num_positions,), "int32"),
+                    ],
+                )
+            )
+        gv = bb.emit_func_output(taken_probs_indices)
+    return gv
+
+
+def _attach_batch_verifier(bb: relax.BlockBuilder):
+    num_nodes = tirx.Var("num_nodes", "int64")
+    nbatch = tirx.Var("nbatch", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    draft_probs = relax.Var("draft_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
+    draft_tokens = relax.Var("draft_tokens", relax.TensorType((num_nodes,), "int32"))
+    model_probs = relax.Var("model_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
+    token_tree_first_child = relax.Var(
+        "token_tree_first_child", relax.TensorType((num_nodes,), "int32")
+    )
+    token_tree_next_sibling = relax.Var(
+        "token_tree_next_sibling", relax.TensorType((num_nodes,), "int32")
+    )
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_nodes,), "float32"))
+    token_tree_parent_ptr = relax.Var("token_tree_parent_ptr", relax.TensorType((nbatch,), "int32"))
+    args = [
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    ]
+    with bb.function("sampler_verify_draft_tokens", args):
+        with bb.dataflow():
+            res = bb.emit_output(
+                relax.call_tir_inplace(
+                    bb.add_func(
+                        batch_spec_verify(vocab_size),
+                        "batch_verify_on_gpu_single_kernel",
+                    ),
+                    args,
+                    inplace_indices=[
+                        args.index(model_probs),
+                        args.index(token_tree_parent_ptr),
+                    ],
+                    out_ty=[
+                        model_probs.ty,
+                        token_tree_parent_ptr.ty,
+                    ],
+                )
+            )
+        gv = bb.emit_func_output(res)
+    return gv

python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py

@@ -0,0 +1,274 @@
+"""A compiler pass that attaches two-stage softmax with temperature."""
+
+from typing import Any, Dict, Optional  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import get_max_num_threads_per_block
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSoftmaxWithTemperature")
+class AttachSoftmaxWithTemperature:
+    """Rewrites one-shot softmax into two-stage softmax."""
+
+    def __init__(
+        self,
+        target: tvm.target.Target,
+        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ) -> None:
+        self.target = target
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _Rewriter(mod, self.target, self.metadata).transform()
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    def __init__(
+        self,
+        mod: IRModule,
+        target: tvm.target.Target,
+        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.target = target
+        self.metadata = metadata
+        self.chunk_size = 4096
+        self.active_vocab_size = self.metadata.get("active_vocab_size") if self.metadata else None
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        batch_size = tirx.Var("batch_size", "int64")
+        vocab_size = tirx.Var("vocab_size", "int64")
+        dtype = "float32"
+        logits = relax.Var("logits", relax.TensorType([batch_size, 1, vocab_size], dtype))
+        temperature = relax.Var("temperature", relax.TensorType([batch_size], dtype))
+        with self.builder_.function("softmax_with_temperature", params=[logits, temperature]):
+            with self.builder_.dataflow():
+                output_struct_info = logits.ty
+                new_shape = relax.ShapeExpr([batch_size, vocab_size])
+                logits = relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    logits,
+                    new_shape,
+                    ty_args=relax.TensorType(new_shape, dtype),
+                )
+                f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(
+                    self.target, self.chunk_size, self.active_vocab_size
+                )
+                chunked_result_struct_info = relax.TensorType(
+                    (batch_size, (vocab_size + self.chunk_size - 1) // self.chunk_size),
+                    "float32",
+                )
+                chunked_results = self.builder_.emit(
+                    relax.call_tir(
+                        self.builder_.add_func(f_chunk_lse, "chunk_lse"),
+                        args=[logits, temperature],
+                        out_ty=[
+                            chunked_result_struct_info,
+                            chunked_result_struct_info,
+                        ],
+                    )
+                )
+                chunked_sum = chunked_results[0]
+                chunked_max = chunked_results[1]
+                softmax = self.builder_.emit(
+                    relax.call_tir(
+                        self.builder_.add_func(f_softmax_with_lse, "softmax_with_chunked_sum"),
+                        args=[logits, temperature, chunked_sum, chunked_max],
+                        out_ty=logits.ty,
+                    )
+                )
+                softmax = self.builder_.emit_output(
+                    relax.call_pure_packed(
+                        "vm.builtin.reshape",
+                        softmax,
+                        output_struct_info.shape,
+                        ty_args=output_struct_info,
+                    )
+                )
+            self.builder_.emit_func_output(softmax)
+        return self.builder_.get()
+
+
+def _get_lse_and_softmax_func(target: tvm.target.Target, chunk_size: int, active_vocab_size: int):
+    # NOTE: A quick note on the softmax implementation.
+    # We once tried to multiply every element by log2e which can be computed
+    # potentially more efficiently on hardware.
+    # However, when the input values are large, multiplying by the factor of log2e
+    # causes numerical issue in float32 dtype.
+    # This leads to the softmax output not summing up to 1.
+    # For numerical stability, we removed the log2e factor and switched back
+    # to the standard log/exp computation.
+
+    # The kernels below handle both the cases of temperature=0 and temperature != 0.
+    # - When temperature is not 0, the first kernel computes the log-sum-exp of
+    # chunks (subtracted by the max value in chunk), and the max values of chunks.
+    # The second kernel merges the log-sum-exp with the maximum values.
+    # - When temperature is 0, the first kernel computes the max value and the counts
+    # of the max value. The second kernel merges the max and counts, and set the
+    # softmax of the maximum values to "max_value / max_count".
+
+    @T.prim_func(s_tir=True)
+    def chunk_lse(
+        var_A: T.handle,
+        var_temperature: T.handle,
+        var_chunked_sum: T.handle,
+        var_chunked_max: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        batch_size = T.int64()
+        vocab_size = T.int64()
+        num_chunks = T.int64()
+        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
+        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
+        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
+        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
+        A_pad = T.sblock_alloc_buffer(
+            (batch_size, num_chunks, T.int64(chunk_size)), dtype="float32"
+        )
+        temp_max = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
+        temp_sum = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
+
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                A_pad[v0, v1, v2] = T.Select(
+                    v1 * T.int64(chunk_size) + v2
+                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                    T.if_then_else(
+                        temperature[v0] > T.float32(1e-5),
+                        A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0],
+                        A[v0, v1 * T.int64(chunk_size) + v2],
+                    ),
+                    T.min_value("float32"),
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("max"):
+                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
+                with T.init():
+                    temp_max[v0, v1] = T.min_value("float32")
+                temp_max[v0, v1] = T.max(temp_max[v0, v1], A_pad[v0, v1, v2])
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("sum_exp"):
+                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
+                with T.init():
+                    temp_sum[v0, v1] = T.float32(0)
+                temp_sum[v0, v1] += T.if_then_else(
+                    v1 * T.int64(chunk_size) + v2
+                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                    T.Select(
+                        temperature[v0] > T.float32(1e-5),
+                        T.exp(A_pad[v0, v1, v2] - temp_max[v0, v1]),
+                        T.cast(A_pad[v0, v1, v2] == temp_max[v0, v1], "float32"),
+                    ),
+                    T.float32(0),
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(1)):
+            with T.sblock("log"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                chunked_sum[v0, v1] = T.Select(
+                    temperature[v0] > T.float32(1e-5),
+                    T.log(temp_sum[v0, v1]),
+                    temp_sum[v0, v1],
+                )
+                chunked_max[v0, v1] = temp_max[v0, v1]
+
+    @T.prim_func(s_tir=True)
+    def softmax_with_chunked_sum(
+        var_A: T.handle,
+        var_temperature: T.handle,
+        var_chunked_sum: T.handle,
+        var_chunked_max: T.handle,
+        var_softmax: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        batch_size = T.int64()
+        vocab_size = T.int64()
+        num_chunks = T.int64()
+        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
+        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
+        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
+        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
+        softmax = T.match_buffer(var_softmax, (batch_size, vocab_size), dtype="float32")
+        temp_max = T.sblock_alloc_buffer((batch_size,), dtype="float32")
+        temp_sum = T.sblock_alloc_buffer((batch_size,), dtype="float32")
+        for l0, l1 in T.grid(batch_size, num_chunks):
+            with T.sblock("max"):
+                v0, v1 = T.axis.remap("SR", [l0, l1])
+                with T.init():
+                    temp_max[v0] = T.min_value("float32")
+                temp_max[v0] = T.max(temp_max[v0], chunked_max[v0, v1])
+        for l0, l1 in T.grid(batch_size, num_chunks):
+            with T.sblock("sum_exp"):
+                v0, v1 = T.axis.remap("SR", [l0, l1])
+                with T.init():
+                    temp_sum[v0] = T.float32(0)
+                temp_sum[v0] += T.Select(
+                    temperature[v0] > T.float32(1e-5),
+                    T.exp(chunked_sum[v0, v1] + chunked_max[v0, v1] - temp_max[v0]),
+                    T.cast(chunked_max[v0, v1] == temp_max[v0], "float32") * chunked_sum[v0, v1],
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("log_pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                if v1 * T.int64(chunk_size) + v2 < vocab_size:
+                    softmax[v0, v1 * T.int64(chunk_size) + v2] = T.Select(
+                        v1 * T.int64(chunk_size) + v2
+                        < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                        T.if_then_else(
+                            temperature[v0] > T.float32(1e-5),
+                            T.exp(
+                                A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0]
+                                - (T.log(temp_sum[v0]) + temp_max[v0])
+                            ),
+                            T.cast(
+                                A[v0, v1 * T.int64(chunk_size) + v2] == temp_max[v0],
+                                "float32",
+                            )
+                            / temp_sum[v0],
+                        ),
+                        T.float32(0),
+                    )
+
+    sch = tvm.s_tir.Schedule(IRModule({"softmax_with_chunked_sum": softmax_with_chunked_sum}))
+
+    def apply_gpu_schedule(target, sch):
+        max_threads = get_max_num_threads_per_block(target)
+        TX = 32
+        TY = max_threads // TX
+        unroll_depth = 64
+
+        sch.work_on("softmax_with_chunked_sum")
+        l0, l1, l2 = sch.get_loops("log_pad")
+        bx = sch.fuse(l0, l1)
+        sch.bind(bx, "blockIdx.x")
+        unroll, ty, tx = sch.split(l2, [None, TY, TX])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.annotate(unroll, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+        sch.annotate(unroll, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        for block_name in ["sum_exp", "max"]:
+            block = sch.get_sblock(block_name)
+            sch.set_scope(block, buffer_index=0, storage_scope="shared")
+            sch.compute_at(block, bx)
+            r_loop = sch.get_loops(block)[-1]
+            r_loop, tx = sch.split(r_loop, [None, TX])
+            sch.reorder(tx, r_loop)
+            sch.bind(tx, "threadIdx.x")
+            sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+            sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        return chunk_lse, sch.mod["softmax_with_chunked_sum"]
+
+    if target.kind.name == "llvm":
+        return chunk_lse, sch.mod["softmax_with_chunked_sum"]
+    return apply_gpu_schedule(target, sch)

python/mlc_llm/compiler_pass/attach_spec_decode_aux_funcs.py

@@ -0,0 +1,123 @@
+"""The pass that attaches logit processor functions to the IRModule."""
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.relax import BlockBuilder, TensorType
+from tvm.script import tirx as T
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSpecDecodeAuxFuncs")
+class AttachSpecDecodeAuxFuncs:
+    """Attach logit processing TIR functions to IRModule."""
+
+    tensor_parallel_shards: int
+
+    def __init__(self, tensor_parallel_shards: int):
+        self.tensor_parallel_shards = tensor_parallel_shards
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        mod = mod.clone()
+        bb = BlockBuilder(mod)
+        bb.add_func(
+            _get_scatter_2d_inplace(dtype="float32", global_symbol="scatter_probs"),
+            "scatter_probs",
+        )
+        bb.add_func(
+            _get_gather_2d_inplace(dtype="float32", global_symbol="gather_probs"),
+            "gather_probs",
+        )
+        if "prefill_to_last_hidden_states" in mod:
+            hidden_states_struct_info = mod["prefill_to_last_hidden_states"].ret_ty.fields[0]
+            dtype = hidden_states_struct_info.dtype
+            _add_gather_hidden_states(bb, self.tensor_parallel_shards, dtype)
+            _add_scatter_hidden_states(bb, self.tensor_parallel_shards, dtype)
+        return bb.finalize()
+
+
+def _get_scatter_2d_inplace(dtype: str, global_symbol: str):
+    @T.prim_func(s_tir=True)
+    def _scatter_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle):
+        T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True})
+        batch_size = T.int32()
+        m = T.int32()
+        n = T.int32()
+        src = T.match_buffer(var_src, (batch_size, n), dtype)
+        indices = T.match_buffer(var_indices, (batch_size,), "int32")
+        dst = T.match_buffer(var_dst, (m, n), dtype)
+        for b, j in T.grid(batch_size, n):
+            with T.sblock("scatter_2d"):
+                vb, vj = T.axis.remap("SS", [b, j])
+                dst[indices[vb], vj] = src[vb, vj]
+
+    return _scatter_2d
+
+
+def _get_gather_2d_inplace(dtype: str, global_symbol: str):
+    @T.prim_func(s_tir=True)
+    def _gather_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle):
+        T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True})
+        batch_size = T.int32()
+        m = T.int32()
+        n = T.int32()
+        src = T.match_buffer(var_src, (m, n), dtype)
+        indices = T.match_buffer(var_indices, (batch_size,), "int32")
+        dst = T.match_buffer(var_dst, (batch_size, n), dtype)
+        for b, j in T.grid(batch_size, n):
+            with T.sblock("gather_2d"):
+                vb, vj = T.axis.remap("SS", [b, j])
+                dst[vb, vj] = src[indices[vb], vj]
+
+    return _gather_2d
+
+
+def _add_scatter_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str):
+    batch_size = tirx.Var("batch_size", "int64")
+    m = tirx.Var("m", "int64")
+    n = tirx.Var("n", "int64")
+    src = relax.Var("src", ty=TensorType([batch_size, n], dtype))
+    indices = relax.Var("indices", ty=TensorType([batch_size], "int32"))
+    dst = relax.Var("dst", ty=TensorType([m, n], dtype))
+    with bb.function("scatter_hidden_states", [src, indices, dst]):
+        with bb.dataflow():
+            if tensor_parallel_shards > 1:
+                indices = relax.op.ccl.broadcast_from_worker0(indices)
+            output = bb.emit_output(
+                relax.op.call_tir_inplace(
+                    bb.add_func(
+                        _get_scatter_2d_inplace(dtype, "_scatter_hidden_states"),
+                        "_scatter_hidden_states",
+                    ),
+                    [src, indices, dst],
+                    2,
+                    dst.ty,
+                )
+            )
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+def _add_gather_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str):
+    batch_size = tirx.Var("batch_size", "int64")
+    m = tirx.Var("m", "int64")
+    n = tirx.Var("n", "int64")
+    src = relax.Var("src", ty=TensorType([m, n], dtype))
+    indices = relax.Var("indices", ty=TensorType([batch_size], "int32"))
+    dst = relax.Var("dst", ty=TensorType([batch_size, n], dtype))
+    with bb.function("gather_hidden_states", [src, indices, dst]):
+        with bb.dataflow():
+            if tensor_parallel_shards > 1:
+                indices = relax.op.ccl.broadcast_from_worker0(indices)
+            output = bb.emit_output(
+                relax.op.call_tir_inplace(
+                    bb.add_func(
+                        _get_gather_2d_inplace(dtype, "_gather_hidden_states"),
+                        "_gather_hidden_states",
+                    ),
+                    [src, indices, dst],
+                    2,
+                    dst.ty,
+                )
+            )
+        gv = bb.emit_func_output(output)
+    return gv

python/mlc_llm/compiler_pass/attach_support_info.py

@@ -0,0 +1,154 @@
+"""A couple of passes that simply supportive information onto the IRModule."""
+
+from math import lcm
+from typing import Any, Dict, List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.ir import Op
+from tvm.relax.expr_functor import PyExprVisitor, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachVariableBounds")
+class AttachVariableBounds:
+    """Attach variable bounds to each Relax function, which primarily helps with memory planning."""
+
+    def __init__(self, variable_bounds: Dict[str, int]):  # noqa: UP006
+        # Specifically for RWKV workloads, which contains -1 max_seq_len
+        self.variable_bounds = {k: v for k, v in variable_bounds.items() if v > 0}
+        self.non_negative_var = ["vocab_size"]
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                mod[g_var] = func.with_attr("tir_var_upper_bound", self.variable_bounds).with_attr(
+                    "tir_non_negative_var", self.non_negative_var
+                )
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachAdditionalPrimFuncs")
+class AttachAdditionalPrimFuncs:
+    """Attach extra TIR PrimFuncs to the IRModule"""
+
+    def __init__(self, functions: Dict[str, tirx.PrimFunc]):  # noqa: UP006
+        self.functions = functions
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for func_name, func in self.functions.items():
+            mod[func_name] = func.with_attr("global_symbol", func_name)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachMemoryPlanAttr")
+class AttachMemoryPlanAttr:
+    """Attach memory planning attribute for dynamic function output planning to Relax functions."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                mod[g_var] = func.with_attr("relax.memory_plan_dynamic_func_output", True)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphCaptureHints")
+class AttachCUDAGraphSymbolicCaptureHints:
+    """Attach CUDA graph capture hints to the IRModule"""
+
+    def __init__(self, hints: Dict[str, List[str]]):  # noqa: UP006
+        self.hints = hints
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            func_name = g_var.name_hint
+            if isinstance(func, relax.Function):
+                if func_name in self.hints:
+                    mod[g_var] = func.with_attr(
+                        "relax.rewrite_cuda_graph.capture_symbolic_vars",
+                        self.hints[func_name],
+                    )
+
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachPipelineParallelStages")
+class AttachPipelineParallelStages:
+    """Attach number of pipeline stages to relax functions."""
+
+    def __init__(self, pipeline_parallel_shards: int):
+        self.pipeline_parallel_shards = pipeline_parallel_shards
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            func_name = g_var.name_hint
+            if not isinstance(func, relax.Function) or func_name not in [
+                "prefill",
+                "decode",
+                "prefill_to_last_hidden_states",
+                "decode_to_last_hidden_states",
+                "batch_prefill",
+                "batch_decode",
+                "batch_verify",
+                "batch_prefill_to_last_hidden_states",
+                "batch_decode_to_last_hidden_states",
+                "batch_verify_to_last_hidden_states",
+            ]:
+                continue
+            mod[g_var] = func.with_attr("pipeline_parallel_stages", self.pipeline_parallel_shards)
+
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSequenceLengthPaddingFactor")
+class AttachSequenceLengthPaddingFactor:
+    """Attach sequence length padding factor to the metadata"""
+
+    def __init__(self, target: tvm.target.Target, metadata: Dict[str, Any]):  # noqa: UP006
+        self.target = target
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+
+        @visitor
+        class _Visitor(PyExprVisitor):
+            def __init__(self, target: tvm.target.Target) -> None:
+                self.padding_factor = 1
+                self.target = target
+                self._op_call_dps_packed = Op.get("relax.call_dps_packed")
+
+            def run(self, mod: IRModule) -> int:
+                """Entry point of the visitor."""
+                # Right now we only need padding for CUDA SM100a architecture.
+                # When the target is SM100a and uses cutlass gemm function,
+                # the sequence length needs to be padded to multiple of 4.
+                if self.target.kind.name != "cuda" or self.target.attrs.get("arch") != "sm_100a":
+                    return 1
+
+                for _, func in mod.functions_items():
+                    if isinstance(func, relax.Function):
+                        self.visit_expr(func)
+                return self.padding_factor
+
+            def visit_call_(self, call: relax.Call) -> None:
+                super().visit_call_(call)
+                if call.op != self._op_call_dps_packed:
+                    return
+                func_name = str(call.args[0].global_symbol)
+                if func_name in [
+                    "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn",
+                    "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn",
+                ]:
+                    # Find the minimum common multiple of padding factor and 4
+                    self.padding_factor = lcm(self.padding_factor, 4)
+
+        # self.metadata["sequence_length_padding"] = True
+        padding_factor = _Visitor(self.target).run(mod)
+        if padding_factor > 1:
+            self.metadata["seqlen_padding_factor"] = padding_factor
+        return mod

python/mlc_llm/compiler_pass/blas_dispatch.py

@@ -0,0 +1,51 @@
+"""A compiler pass that dispatches patterns to CUBLAS."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.backend import get_patterns_with_prefix
+
+try:
+    import tvm.relax.backend.cuda.cublas as _cublas  # noqa: F401
+    import tvm.relax.backend.rocm.hipblas as _hipblas  # noqa: F401
+except ImportError:
+    # Note: legacy path of cublas/hipblas for backward compatibility
+    pass
+
+
+@tvm.transform.module_pass(opt_level=0, name="BLASDispatch")
+class BLASDispatch:
+    """A compiler pass that dispatches patterns to cuBLAS/hipBLAS."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        if target.kind.name == "cuda":
+            self.has_blas = tvm.get_global_func("relax.ext.cublas", True)
+            if not self.has_blas:
+                raise Exception("cuBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("cublas")
+        elif target.kind.name == "rocm":
+            self.has_blas = tvm.get_global_func("relax.ext.hipblas", True)
+            if not self.has_blas:
+                raise Exception("hipBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("hipblas")
+        else:
+            raise Exception(f"Unsupported target {target.kind.name} for BLAS dispatch.")
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        model_names = [
+            gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function)
+        ]
+        # exclude single batch decode
+        model_names = [name for name in model_names if "batch" in name or "decode" not in name]
+        mod = tvm.transform.Sequential(
+            [
+                relax.transform.FuseOpsByPattern(
+                    self.patterns,
+                    bind_constants=False,
+                    annotate_codegen=True,
+                    entry_functions=model_names,
+                ),
+                relax.transform.RunCodegen({}, entry_functions=model_names),
+            ]
+        )(mod)
+        return mod

python/mlc_llm/compiler_pass/clean_up_tir_attrs.py

@@ -0,0 +1,31 @@
+"""A compiler pass that cleans up undesired TIR attrs."""
+
+from typing import List  # noqa: UP035
+
+import tvm
+from tvm.ir.module import IRModule
+
+
+@tvm.transform.module_pass(opt_level=0, name="CleanUpTIRAttrs")
+class CleanUpTIRAttrs:
+    """A compiler pass that cleans up undesired TIR attrs."""
+
+    def __init__(self, attrs: List[str]):  # noqa: UP006
+        self.attrs = attrs
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for g_var, func in mod.functions_items():
+            changed = False
+            for attr in self.attrs:
+                if func.attrs is not None and attr in func.attrs:
+                    func = func.without_attr(attr)
+                    changed = True
+                    break
+            if changed:
+                mod[g_var] = func
+        return mod

python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py

@@ -0,0 +1,243 @@
+"""A pass that rewrites KV cache creation functions in IRModule."""
+
+import json
+from typing import Any, Dict, List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.frontend.nn.llm import kv_cache
+from tvm.relax.frontend.nn.llm.kv_cache import RopeMode
+
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+def extract_creation_args(func: relax.Function) -> Dict[str, Any]:  # noqa: UP006
+    """Extract the KV cache creation args from the given generic creation func."""
+    assert isinstance(func.body, relax.SeqExpr)
+    assert len(func.body.blocks) == 1
+    assert isinstance(func.body.blocks[0], relax.DataflowBlock)
+    assert isinstance(func.body.blocks[0].bindings[0], relax.VarBinding)
+    assert isinstance(func.body.blocks[0].bindings[0].value, relax.Call)
+    assert func.body.blocks[0].bindings[0].value.op == tvm.ir.Op.get("relax.call_pure_packed")
+    call_args = func.body.blocks[0].bindings[0].value.args
+    assert isinstance(call_args[0], relax.ExternFunc)
+    assert call_args[0].global_symbol == "mlc.create_paged_kv_cache_generic"
+    args = call_args[1:]
+    assert len(args) == 18
+    assert isinstance(args[0], (relax.StringImm, relax.Tuple))
+    # Check if attn_kind is a single value or a list with length of hidden layers
+    if isinstance(args[0], relax.StringImm):
+        assert args[0].value in ["mha", "mla"]
+        attn_kind = args[0].value
+    else:
+        assert len(args[0].fields) == args[3].value
+        for i, attention_type in enumerate(args[0].fields):
+            assert isinstance(attention_type, relax.StringImm)
+            assert attention_type.value in ["mha", "mla", "mha_sliding"]
+        attn_kind = [args[0].fields[i].value for i in range(len(args[0]))]
+    assert isinstance(args[1], relax.ShapeExpr)
+    assert len(args[1].values) == 5
+    assert isinstance(args[2], relax.ShapeExpr)
+    for i in range(3, 18):
+        if i in [13, 14, 17]:
+            continue
+        # PrimValue wrappers were phased out of Relax: scalar args are now bare
+        # tirx PrimExprs (IntImm/FloatImm) directly.
+        assert isinstance(args[i], (tvm.tirx.IntImm, tvm.tirx.FloatImm)), (
+            f"args[{i}] is {type(args[i])}"
+        )
+    assert isinstance(args[13], relax.StringImm)
+    assert isinstance(args[16], (relax.Constant, tvm.tirx.IntImm, tvm.tirx.FloatImm))
+    assert isinstance(args[17], relax.DataTypeImm)
+
+    return {
+        "attn_kind": attn_kind,
+        "max_batch_size": args[1].values[0],
+        "max_total_seq_len": args[1].values[1],
+        "prefill_chunk_size": args[1].values[2],
+        "page_size": args[1].values[3],
+        "support_sliding_window": args[1].values[4],
+        "layer_partition": args[2],
+        "num_hidden_layers": args[3].value,
+        "num_attention_heads": args[4].value,
+        "num_key_value_heads": args[5].value,
+        "qk_head_dim": args[6].value,
+        "v_head_dim": args[7].value,
+        "mla_original_qk_head_dim": args[8].value,
+        "mla_original_v_head_dim": args[9].value,
+        "rope_mode": args[10].value,
+        "rope_scale": args[11].value,
+        "rope_theta": args[12].value,
+        "rope_scaling": json.loads(args[13].value),
+        "rope_ext_factors": args[14],
+        "rotary_dim": args[15].value,
+        "enable_disaggregation": bool(args[16].value),
+        "dtype": args[17].value,
+    }
+
+
+@tvm.transform.module_pass(opt_level=0, name="DispatchKVCacheCreation")
+class DispatchKVCacheCreation:
+    """Rewrite KV cache creation functions to IRModule."""
+
+    def __init__(
+        self,
+        target: tvm.target.Target,
+        flashinfer: bool,
+        metadata: Dict[str, Any],  # noqa: UP006
+    ) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            The target of the model compilation.
+
+        flashinfer : bool
+            A boolean indicating if flashinfer is enabled.
+
+        metadata : Dict[str, Any]
+            The model's metadata for KV cache creation.
+            Note that the metadata will be updated in this pass -- the
+            KV cache metadata will be attached.
+        """
+        self.target = target
+        self.flashinfer = flashinfer
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        func_dict = {}
+        creation_func = None
+        for g_var, func in mod.functions_items():
+            # Try to find the `create_paged_kv_cache` func.
+            if g_var.name_hint == "create_paged_kv_cache":
+                creation_func = func
+            else:
+                func_dict[g_var] = func
+
+        if creation_func is None:
+            return mod
+
+        new_mod = IRModule(func_dict)
+        if mod.attrs is not None:
+            new_mod = new_mod.with_attrs(mod.attrs)
+
+        kwargs = extract_creation_args(creation_func)
+        self.attach_kv_cache_metadata(kwargs)
+
+        bb = relax.BlockBuilder(new_mod)
+        extern_mods = []
+        extern_mods += self.create_tir_paged_kv_cache(bb, kwargs)
+        extern_mods += self.create_flashinfer_paged_kv_cache(bb, kwargs)
+
+        mod = bb.finalize()
+        mod_attrs = dict(mod.attrs) if mod.attrs else {}
+        mod = mod.with_attr("external_mods", mod_attrs.get("external_mods", []) + extern_mods)
+        return mod
+
+    def attach_kv_cache_metadata(self, kwargs: Dict[str, Any]):  # noqa: UP006
+        """Attach the KV cache metadata to model metadata."""
+        self.metadata["kv_cache"] = {
+            "num_hidden_layers": kwargs["num_hidden_layers"],
+            "num_attention_heads": kwargs["num_attention_heads"],
+            "num_key_value_heads": kwargs["num_key_value_heads"],
+            "head_dim": kwargs["qk_head_dim"],
+        }
+
+    def create_tir_paged_kv_cache(
+        self,
+        bb: relax.BlockBuilder,
+        kwargs: Dict[str, Any],  # noqa: UP006
+    ) -> List[tvm.runtime.Module]:  # noqa: UP006
+        """Create the TIR-based PagedKVCache"""
+        max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
+        max_total_seq_len = relax.Var(
+            "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
+        )
+        prefill_chunk_size = relax.Var(
+            "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
+        )
+        page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
+        support_sliding_window = relax.Var(
+            "support_sliding_window_",
+            relax.ShapeType([kwargs["support_sliding_window"]]),
+        )
+
+        # Ensure 'enable_disaggregation' is optional
+        enable_disaggregation = kwargs.pop("enable_disaggregation", False)
+        kwargs["enable_disaggregation"] = enable_disaggregation
+
+        with bb.function(
+            name="create_tir_paged_kv_cache",
+            params=[
+                max_batch_size,
+                max_total_seq_len,
+                prefill_chunk_size,
+                page_size,
+                support_sliding_window,
+            ],
+        ):
+            cache = kv_cache.TIRPagedKVCache(target=self.target, **kwargs)
+            bb.emit_func_output(cache._expr)
+
+        return cache.extern_mods
+
+    def create_flashinfer_paged_kv_cache(
+        self,
+        bb: relax.BlockBuilder,
+        kwargs: Dict[str, Any],  # noqa: UP006
+    ) -> List[tvm.runtime.Module]:  # noqa: UP006
+        """Create the FlashInfer-based PagedKVCache"""
+        # Filter the cases which FlashInfer does not support.
+        if (
+            not self.flashinfer
+            or self.target.kind.name != "cuda"
+            or str(kwargs["dtype"]) not in ["float16", "bfloat16"]
+            or (
+                kwargs["rope_mode"] == RopeMode.INLINE
+                and (
+                    kwargs["rotary_dim"] != kwargs["qk_head_dim"]
+                    or kwargs["qk_head_dim"] != kwargs["v_head_dim"]
+                )
+            )
+        ):
+            return []
+
+        max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
+        max_total_seq_len = relax.Var(
+            "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
+        )
+        prefill_chunk_size = relax.Var(
+            "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
+        )
+        page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
+        support_sliding_window = relax.Var(
+            "support_sliding_window_",
+            relax.ShapeType([kwargs["support_sliding_window"]]),
+        )
+
+        try:
+            with bb.function(
+                name="create_flashinfer_paged_kv_cache",
+                params=[
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ],
+            ):
+                cache = kv_cache.FlashInferPagedKVCache(target=self.target, **kwargs)
+                bb.emit_func_output(cache._expr)
+        except Exception as e:
+            logger.info(
+                "Error caught when creating FlashInfer PagedKVCache: %s\n"
+                "The model will fallback to TIR-based KV cache.",
+                e,
+            )
+            return []
+
+        return cache.extern_mods

python/mlc_llm/compiler_pass/dispatch_triton_kernel.py

@@ -0,0 +1,176 @@
+"""A pass that dispatch generic calls of triton kernels to specific kernel implementations."""
+
+from typing import List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+from mlc_llm.op.triton import (
+    get_tir_w8a8_block_fp8_group_matmul,
+    get_tir_w8a8_block_fp8_matmul,
+)
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    def __init__(self, mod: IRModule, target: tvm.target.Target) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.target = target
+        self.extern_mods: List[tvm.runtime.Module] = []  # noqa: UP006
+
+    def transform(self) -> tvm.IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function):
+                continue
+            new_func = self.visit_expr(func)
+            # new_func = remove_all_unused(new_func)
+            self.builder_.update_func(g_var, new_func)
+
+        mod = self.builder_.finalize()
+        mod_attrs = dict(mod.attrs) if mod.attrs else {}
+        mod = mod.with_attr(
+            "external_mods", list(mod_attrs.get("external_mods", [])) + self.extern_mods
+        )
+        return mod
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        call = super().visit_call_(call)
+
+        if (
+            call.op != tvm.ir.Op.get("relax.call_dps_packed")
+            or not isinstance(call.args[0], relax.ExternFunc)
+            or not str(call.args[0].global_symbol).startswith("mlc.triton.")
+        ):
+            return call
+
+        global_symbol = str(call.args[0].global_symbol)
+        assert isinstance(call.args[1], relax.Tuple)
+        if global_symbol == "mlc.triton.w8a8_block_fp8_matmul":
+            return self.w8a8_block_fp8_matmul(call.args[1].fields, call.ty)
+        if global_symbol == "mlc.triton.w8a8_block_fp8_group_matmul":
+            return self.w8a8_block_fp8_group_matmul(call.args[1].fields, call.ty)
+        raise ValueError(f"Unknown mlc.triton kernel identifier: {global_symbol}")
+
+    def w8a8_block_fp8_matmul(
+        self,
+        args: List[relax.Expr],  # noqa: UP006
+        out_ty: relax.Type,
+    ) -> relax.Expr:
+        """Emit the w8a8_block_fp8_matmul triton kernel."""
+        assert len(args) == 16
+        x, weight, x_scale, weight_scale = args[:4]
+        (
+            N,
+            K,
+            block_n,
+            block_k,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+        ) = [arg.value.value for arg in args[4:14]]
+        in_dtype, out_dtype = str(args[14].value), str(args[15].value)
+
+        prim_func, func_name = get_tir_w8a8_block_fp8_matmul(
+            N,
+            K,
+            block_n,
+            block_k,
+            in_dtype,
+            out_dtype,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            self.extern_mods,
+        )
+        if prim_func is None:
+            # The TIR function is already in the IRModule
+            gv = self.builder_.get().get_global_var(func_name)
+        else:
+            # Add the TIR function to the IRModule
+            gv = self.builder_.add_func(prim_func, func_name)
+
+        return relax.call_tir(gv, [x, weight, x_scale, weight_scale], out_ty=out_ty)
+
+    def w8a8_block_fp8_group_matmul(
+        self,
+        args: List[relax.Expr],  # noqa: UP006
+        out_ty: relax.Type,
+    ) -> relax.Expr:
+        """Emit the w8a8_block_fp8_group_matmul triton kernel."""
+        assert len(args) == 19
+        x, weight, x_scale, weight_scale, expert_ids, indptr = args[:6]
+        (
+            N,
+            K,
+            num_experts,
+            block_n,
+            block_k,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+        ) = [arg.value.value for arg in args[6:17]]
+        in_dtype, out_dtype = str(args[17].value), str(args[18].value)
+
+        prim_func, func_name = get_tir_w8a8_block_fp8_group_matmul(
+            N,
+            K,
+            num_experts,
+            block_n,
+            block_k,
+            in_dtype,
+            out_dtype,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            self.extern_mods,
+        )
+        if prim_func is None:
+            # The TIR function is already in the IRModule
+            gv = self.builder_.get().get_global_var(func_name)
+        else:
+            # Add the TIR function to the IRModule
+            gv = self.builder_.add_func(prim_func, func_name)
+
+        return relax.call_tir(
+            gv,
+            [x, weight, x_scale, weight_scale, expert_ids, indptr],
+            out_ty=out_ty,
+        )
+
+
+@tvm.transform.module_pass(opt_level=0, name="DispatchTritonKernel")
+class DispatchTritonKernel:
+    """Rewrite KV cache creation functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        """
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        if self.target.kind.name != "cuda":
+            return mod
+
+        return _Rewriter(mod, self.target).transform()

python/mlc_llm/compiler_pass/estimate_memory_usage.py

@@ -0,0 +1,87 @@
+"""Memory usage estimation analysis function for Relax functions."""
+
+import json
+from typing import Any, Dict  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir import IRModule, Op
+from tvm.relax.expr_functor import PyExprVisitor, visitor
+
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachMetadata")
+class AttachMetadataWithMemoryUsage:
+    """Attach a Relax function that returns metadata in a JSON string"""
+
+    def __init__(self, metadata: Dict[str, Any]):  # noqa: UP006
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+
+        func_name = "_metadata"
+
+        def _emit_metadata(metadata):
+            bb = relax.BlockBuilder()
+            with bb.function(func_name, params=[]):
+                bb.emit_func_output(relax.StringImm(json.dumps(metadata)))
+            return bb.finalize()[func_name]
+
+        self.metadata["memory_usage"] = _MemoryEstimator().run(mod)
+        mod[func_name] = _emit_metadata(self.metadata)
+        return mod
+
+
+@visitor
+class _MemoryEstimator(PyExprVisitor):
+    """The IR visitor which estimates the memory usage of each Relax function."""
+
+    def __init__(self) -> None:
+        self.planned_alloc_mem = 0
+        self.planned_mem_num = 0
+        self._op_alloc_tensor = Op.get("relax.builtin.alloc_tensor")
+        self._op_alloc_storage = Op.get("relax.memory.alloc_storage")
+
+    def run(self, mod: IRModule) -> Dict[str, int]:  # noqa: UP006
+        """Entry point of the visitor."""
+        result: Dict[str, int] = {}  # noqa: UP006
+        for global_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                self.planned_alloc_mem = 0
+                self.planned_mem_num = 0
+                self.visit_expr(func)
+                result[global_var.name_hint] = self.planned_alloc_mem
+                logger.info(
+                    "[Memory usage] Function `%s`: %.2f MB",
+                    global_var.name_hint,
+                    self.planned_alloc_mem / 1024 / 1024,
+                )
+        return result
+
+    def visit_call_(self, call: relax.Call) -> None:
+        if call.op == self._op_alloc_tensor:
+            self._builtin_tensor_alloc(shape=call.args[0], dtype_str=call.args[1].value)
+        elif call.op == self._op_alloc_storage:
+            self._storage_alloc(size=call.args[0])
+        super().visit_call_(call)
+
+    def _builtin_tensor_alloc(self, shape: relax.Expr, dtype_str: str) -> None:
+        assert isinstance(shape, relax.ShapeExpr)
+        size = 1
+        for dim_len in shape.values:
+            if not isinstance(dim_len, tvm.tirx.IntImm):
+                return
+            size *= dim_len.value
+        dtype = tvm.DataType(dtype_str)
+        self.planned_mem_num += 1
+        self.planned_alloc_mem += size * ((dtype.bits + 7) // 8) * dtype.lanes
+
+    def _storage_alloc(self, size: relax.Expr) -> None:
+        assert isinstance(size, relax.ShapeExpr)
+        if isinstance(size.values[0], tirx.IntImm):
+            self.planned_mem_num += 1
+            self.planned_alloc_mem += size.values[0].value

python/mlc_llm/compiler_pass/fuse_add_norm.py

@@ -0,0 +1,238 @@
+"""A compiler pass that fuses add + rms_norm."""
+
+from typing import Optional
+
+import tvm
+from tvm import relax
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import get_max_num_threads_per_block
+
+
+def _get_add_rms_norm_decode(hidden_size: int, eps: float, TX: int, in_dtype: str):
+    if in_dtype not in ("float16", "bfloat16"):
+        raise ValueError(f"Unsupported data type: {in_dtype}")
+    inv_hidden_size = T.float32(1.0 / float(hidden_size))
+    eps = T.float32(eps)
+    add_local_size = hidden_size // TX
+
+    @T.prim_func(private=True, s_tir=True)
+    def decode_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        batch_size = T.int32()
+        A = T.match_buffer(pA, (batch_size, 1, hidden_size), in_dtype)
+        B = T.match_buffer(pB, (batch_size, 1, hidden_size), in_dtype)
+        C = T.match_buffer(pC, (hidden_size,), in_dtype)
+        out = T.match_buffer(pO, (batch_size, 1, hidden_size), in_dtype)
+        add = T.match_buffer(pAdd, (batch_size, 1, hidden_size), in_dtype)
+        add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
+        sum_shared = T.sblock_alloc_buffer((batch_size, 1), scope="shared")
+        sum_local = T.sblock_alloc_buffer((TX, batch_size, 1), scope="local")
+        for v_bx in T.thread_binding(batch_size, thread="blockIdx.x"):
+            for v_tx in T.thread_binding(
+                TX,
+                thread="threadIdx.x",
+                annotations={
+                    "pragma_auto_unroll_max_step": 256,
+                    "pragma_unroll_explicit": 1,
+                },
+            ):
+                for i in range(add_local_size):
+                    with T.sblock("T_add"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx)
+                        add_local[h // TX] = A[bx, 0, h] + B[bx, 0, h]
+                    with T.sblock("T_write_back"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        v_ax1 = T.axis.spatial(1, 0)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx)
+                        add[bx, v_ax1, h] = add_local[h // TX]
+                with T.sblock("T_multiply_red_rf_init"):
+                    tx, bx = T.axis.remap("SS", [v_tx, v_bx])
+                    sum_local[tx, bx, 0] = T.float32(0)
+                for v_i, _j in T.grid(add_local_size, 1):
+                    with T.sblock("T_multiply_red_rf_update"):
+                        tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
+                        sum_local[tx, bx, 0] += T.float32(add_local[i]) * T.float32(add_local[i])
+            for _j in range(1):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_multiply_red"):
+                        tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
+                        T.reads(sum_local[tx, bx, 0])
+                        T.writes(sum_shared[bx, 0])
+                        with T.init():
+                            sum_shared[bx, 0] = T.float32(0)
+                        sum_shared[bx, 0] += sum_local[tx, bx, 0]
+            for i in range(add_local_size):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_cast_2"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx_2)
+                        out[bx, 0, h] = T.cast(
+                            T.rsqrt(sum_shared[bx, 0] * inv_hidden_size + eps)
+                            * T.float32(add_local[h // TX])
+                            * T.float32(C[h]),
+                            dtype=in_dtype,
+                        )
+
+    return decode_add_rms
+
+
+def _get_add_rms_norm_prefill(hidden_size: int, eps: float, TX: int, in_dtype: str):
+    if in_dtype not in ("float16", "bfloat16"):
+        raise ValueError(f"Unsupported data type: {in_dtype}")
+    inv_hidden_size = T.float32(1.0 / float(hidden_size))
+    eps = T.float32(eps)
+    add_local_size = hidden_size // TX
+
+    @T.prim_func(private=True, s_tir=True)
+    def prefill_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        seq_len = T.int32()
+        A = T.match_buffer(pA, (1, seq_len, hidden_size), in_dtype)
+        B = T.match_buffer(pB, (1, seq_len, hidden_size), in_dtype)
+        C = T.match_buffer(pC, (hidden_size,), in_dtype)
+        out = T.match_buffer(pO, (1, seq_len, hidden_size), in_dtype)
+        add = T.match_buffer(pAdd, (1, seq_len, hidden_size), in_dtype)
+        add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
+        sum_shared = T.sblock_alloc_buffer((1, seq_len), scope="shared")
+        sum_local = T.sblock_alloc_buffer((TX, 1, seq_len), scope="local")
+        for v_bx in T.thread_binding(seq_len, thread="blockIdx.x"):
+            for v_tx in T.thread_binding(
+                TX,
+                thread="threadIdx.x",
+                annotations={
+                    "pragma_auto_unroll_max_step": 256,
+                    "pragma_unroll_explicit": 1,
+                },
+            ):
+                for v_i in range(add_local_size):
+                    with T.sblock("T_add"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
+                        add_local[h // TX] = A[0, bx, h] + B[0, bx, h]
+                    with T.sblock("T_write_back"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
+                        add[0, bx, h] = add_local[h // TX]
+                with T.sblock("T_multiply_red_rf_init"):
+                    tx, bx = T.axis.remap("SS", [v_tx, v_bx])
+                    sum_local[tx, 0, bx] = T.float32(0)
+                for v_i, _j in T.grid(add_local_size, 1):
+                    with T.sblock("T_multiply_red_rf_update"):
+                        tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
+                        sum_local[tx, 0, bx] += T.float32(add_local[i]) * T.float32(add_local[i])
+            for _j in range(1):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_multiply_red"):
+                        tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
+                        with T.init():
+                            sum_shared[0, bx] = T.float32(0)
+                        sum_shared[0, bx] = sum_shared[0, bx] + sum_local[tx, 0, bx]
+            for v_i in range(add_local_size):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_cast_2"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        v1 = T.axis.spatial(hidden_size, v_i * TX + v_tx_2)
+                        out[0, bx, v1] = T.cast(
+                            T.rsqrt(sum_shared[0, bx] * inv_hidden_size + eps)
+                            * T.float32(add_local[v1 // TX])
+                            * T.float32(C[v1]),
+                            dtype=in_dtype,
+                        )
+
+    return prefill_add_rms
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseAddRMSNorm")
+class FuseAddRMSNorm:
+    """A compiler pass that fuses add + rms_norm."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            Target device.
+        """
+        self.target = target
+
+    def transform_module(self, mod: tvm.IRModule, _ctx: tvm.transform.PassContext) -> tvm.IRModule:
+        """IRModule-level transformation."""
+        return _FuseAddRMSNormRewriter(mod.clone(), self.target).transform()
+
+
+@mutator
+class _FuseAddRMSNormRewriter(PyExprMutator):
+    def __init__(self, mod: tvm.IRModule, target: tvm.target.Target):
+        super().__init__(mod)
+        self.mod = mod
+        self.prefill_norm_gv: Optional[tvm.ir.GlobalVar] = None
+        self.decode_norm_gv: Optional[tvm.ir.GlobalVar] = None
+        self.TX = min(1024, get_max_num_threads_per_block(target))
+
+    def transform(self) -> tvm.IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function):
+                continue
+            new_func = self.visit_expr(func)
+            new_func = remove_all_unused(new_func)
+            self.builder_.update_func(g_var, new_func)
+        return self.builder_.finalize()
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        call = super().visit_call_(call)
+
+        # Match the "rms_norm(add(x1, x2), w)" pattern
+        if call.op != tvm.ir.Op.get("relax.nn.rms_norm") or call.ty.dtype not in [
+            "bfloat16",
+            "float16",
+        ]:
+            return call
+        assert len(call.args) == 2
+        weight = call.args[1]
+        eps = call.attrs.epsilon
+        assert isinstance(call.args[0], relax.Var)
+        y = self.lookup_binding(call.args[0])
+        if not isinstance(y, relax.Call) or y.op != tvm.ir.Op.get("relax.add"):
+            return call
+        assert len(y.args) == 2
+        x1 = y.args[0]
+        x2 = y.args[1]
+        # Extra check
+        n, _, h = x1.ty.shape
+        h = int(h)
+        if h % self.TX != 0:
+            return call
+
+        is_prefill = n == 1
+        func_gv = self.prefill_norm_gv if is_prefill else self.decode_norm_gv
+        if func_gv is None:
+            if is_prefill:
+                func_gv = self.builder_.add_func(
+                    _get_add_rms_norm_prefill(h, eps, self.TX, call.ty.dtype),
+                    "fuse_add_norm_prefill",
+                )
+                self.prefill_norm_gv = func_gv
+            else:
+                func_gv = self.builder_.add_func(
+                    _get_add_rms_norm_decode(h, eps, self.TX, call.ty.dtype),
+                    "fuse_add_norm_decode",
+                )
+                self.decode_norm_gv = func_gv
+
+        tuple_output = self.builder_.emit(
+            relax.call_tir(
+                func_gv,
+                [x1, x2, weight],
+                out_ty=[x1.ty, x2.ty],
+            )
+        )
+        new_o = relax.TupleGetItem(tuple_output, 0)
+        new_y = self.builder_.emit(relax.TupleGetItem(tuple_output, 1))
+        self.set_var_remap(call.args[0], new_y)
+        return new_o

python/mlc_llm/compiler_pass/fuse_dequantize_matmul_ewise.py

@@ -0,0 +1,85 @@
+"""A compiler pass that fuses dequantize + matmul + elementwise."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeMatmulEwise")
+class FuseDequantizeMatmulEwise:
+    """A compiler pass that fuses dequantize + matmul + elementwise."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        seq = []
+        for n_aux_tensor in [0, 1, 2, 3, 4]:
+            for match_ewise in [0, 1, 2, 3, 6]:
+                if match_ewise == 6 and n_aux_tensor != 4:
+                    continue
+                seq.append(
+                    relax.transform.FuseOpsByPattern(
+                        [
+                            (
+                                "dequantize_matmul",
+                                *_pattern(match_ewise, n_aux_tensor),
+                            )
+                        ]
+                    )
+                )
+        seq.append(relax.transform.FuseTIR())
+        return tvm.transform.Sequential(seq)(mod)
+
+
+def _pattern(match_ewise: int, n_aux_tensor: int):
+    w_scaled = wildcard()
+    x = wildcard()
+    w = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([w_scaled] + [wildcard() for _ in range(n_aux_tensor)]),
+        add_constraint=False,
+    )
+    matmul = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([x, w] + [wildcard() for _ in range(match_ewise)]),
+        add_constraint=False,
+    )
+    annotations = {
+        "w_scaled": w_scaled,
+        "x": x,
+        "w": w,
+        "matmul": matmul,
+    }
+
+    def _check_decoding(ctx: relax.transform.PatternCheckContext) -> bool:
+        call = ctx.annotated_expr["w"]
+        if not isinstance(call, relax.Call):
+            return False
+        g_var = call.args[0]
+        if not isinstance(g_var, relax.GlobalVar):
+            return False
+        return g_var.name_hint.startswith("dequantize") or g_var.name_hint.startswith(
+            "fused_dequantize"
+        )
+
+    def _check_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
+        call = ctx.annotated_expr["matmul"]
+        if not isinstance(call, relax.Call):
+            return False
+        g_var = call.args[0]
+        if not isinstance(g_var, relax.GlobalVar):
+            return False
+        return (
+            g_var.name_hint.startswith("matmul")
+            or g_var.name_hint.startswith("fused_matmul")
+            or g_var.name_hint.startswith("NT_matmul")
+            or g_var.name_hint.startswith("fused_NT_matmul")
+        )
+
+    def _check(ctx: relax.transform.PatternCheckContext) -> bool:
+        return _check_decoding(ctx) and _check_matmul(ctx)
+
+    return matmul, annotations, _check

python/mlc_llm/compiler_pass/fuse_dequantize_take.py

@@ -0,0 +1,91 @@
+"""A compiler pass that fuses dequantize + take."""
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.relax.dpl.pattern import (
+    GlobalVarPattern,
+    TuplePattern,
+    is_const,
+    is_op,
+    wildcard,
+)
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTake")
+class FuseDequantizeTake:
+    """A compiler pass that fuses dequantize + take."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        seq = []
+        for n_aux_tensor in [2, 3]:
+            for match_tir_vars in [False, True]:
+                seq.append(
+                    relax.transform.FuseOpsByPattern(
+                        [
+                            (
+                                "dequantize_take",
+                                *_pattern(n_aux_tensor, match_tir_vars),
+                            )
+                        ]
+                    )
+                )
+        seq.append(relax.transform.FuseTIR())
+        mod = tvm.transform.Sequential(seq)(mod)
+        for g_var, func in mod.functions_items():
+            name = g_var.name_hint
+            if isinstance(func, tirx.PrimFunc) and (
+                ("fused_dequantize" in name) and ("take" in name)
+            ):
+                sch_mod = tvm.IRModule({"main": func})
+                sch_mod = tirx.transform.ForceNarrowIndexToInt32()(sch_mod)
+                sch = tvm.s_tir.Schedule(sch_mod)
+                sch.compute_inline("dequantize")
+                mod[g_var] = sch.mod["main"]
+        return mod
+
+
+def _pattern(n_aux_tensor: int, match_tir_vars: bool):
+    dequantize = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([wildcard() for _ in range(n_aux_tensor)]),
+        add_constraint=False,
+    )
+    indices = ~is_const()
+    if match_tir_vars:
+        call_tir_args_take = [
+            GlobalVarPattern(),
+            TuplePattern([dequantize, indices]),
+            wildcard(),
+        ]
+    else:
+        call_tir_args_take = [
+            GlobalVarPattern(),
+            TuplePattern([dequantize, indices]),
+        ]
+    take = is_op("relax.call_tir")(
+        *call_tir_args_take,
+        add_constraint=False,
+    )
+    annotations = {
+        "take": take,
+        "dequantize": dequantize,
+        "indices": indices,
+    }
+
+    def _check(ctx: relax.transform.PatternCheckContext) -> bool:
+        take = ctx.annotated_expr["take"]
+        dequantize = ctx.annotated_expr["dequantize"]
+        if not isinstance(dequantize, relax.Call):
+            return False
+        if not isinstance(take.args[0], relax.GlobalVar) or not isinstance(
+            dequantize.args[0], relax.GlobalVar
+        ):
+            return False
+        return "take" in take.args[0].name_hint and "dequantize" in dequantize.args[0].name_hint
+
+    return take, annotations, _check

python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py

@@ -0,0 +1,107 @@
+"""A compiler pass that fuses transpose + dequantize."""
+
+import tvm
+from tvm import relax, s_tir, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose")
+class FuseDequantizeTranspose:
+    """A compiler pass that fuses transpose + dequantize."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _DequantizeTransposeFuser(mod).transform()
+
+
+@mutator
+class _DequantizeTransposeFuser(PyExprMutator):
+    def __init__(
+        self,
+        mod: IRModule,
+    ):
+        super().__init__(mod)
+        self.mod = mod
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        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()
+
+    def visit_call_(
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        call = self.visit_expr_post_order(call)
+        if call.op != tvm.ir.Op.get("relax.matmul"):
+            return call
+        # Do not fuse dequantize-transpose for GeMM
+        if (
+            call.args[0].ty.ndim < 2
+            or not isinstance(call.args[0].ty.shape[-2], tirx.IntImm)
+            or call.args[0].ty.shape[-2].value != 1
+        ):
+            return call
+
+        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].ty.ndim != 2
+            or matmul_rhs.attrs.axes is not None
+        ):
+            return call
+
+        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.ty, relax.TensorType)
+        ):
+            return call
+
+        dequantize_tir_func = self.mod[transpose_input.args[0]]
+        assert isinstance(dequantize_tir_func, tirx.PrimFunc)
+        if (
+            len(dequantize_tir_func.body.block.alloc_buffers) != 1
+            or not isinstance(dequantize_tir_func.body.block.body, tirx.SeqStmt)
+            or len(dequantize_tir_func.body.block.body) != 2
+            or not isinstance(dequantize_tir_func.body.block.body[1], tirx.For)
+            or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tirx.SBlockRealize)
+            or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose"
+        ):
+            return call
+
+        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]
+        new_func = tirx.PrimFunc(
+            params=new_func_buffers,
+            body=tirx.SBlockRealize(
+                iter_values=[],
+                predicate=True,
+                block=tirx.SBlock(
+                    iter_vars=[],
+                    reads=[],
+                    writes=[],
+                    name_hint="root",
+                    body=dequantize_tir_func.body.block.body[0],
+                ),
+            ),
+        )
+        # Call `renew_defs` for deep-copy to avoid IR node duplication in
+        # different PrimFuncs of an IRModule.
+        new_func = s_tir.renew_defs(new_func)
+        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_ty=matmul_rhs.ty)
+        )
+        return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype)

python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py

@@ -0,0 +1,331 @@
+"""A compiler pass that fuses dequantize matmul + epilogue."""
+
+import operator
+from functools import reduce
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.dpl import rewrite_call
+from tvm.relax.dpl.pattern import is_op, wildcard
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeEpilogue")
+class FuseFTDequantizeEpilogue:
+    """A compiler pass that fuses FasterTransformer dequantize matmul + epilogue."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for gv, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = fuse_bias(func)
+                func = fuse_activation(func)
+                func = fuse_residual_binary(func)
+                func = fuse_residual_unary(func)
+                mod[gv] = func
+        return mod
+
+
+def fuse_bias(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.add` into fastertransformer.gemm_fp16_int as bias:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int", ...)
+    lv2 = relax.add(lv1, bias)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., bias, ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    bias = wildcard()
+    pattern = is_op("relax.add")(decode_matmul, bias) | is_op("relax.add")(bias, decode_matmul)
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int":
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 8
+            if not args_list[3].value == "identity":
+                # bias cannot be fused after activation
+                return expr
+            matched_bias = match[bias]
+            bias_stride = (
+                matched_bias.ty.shape[-1]
+                if bias
+                and not reduce(operator.mul, matched_bias.ty.shape, 1) == matched_bias.ty.shape[-1]
+                else 0
+            )
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    matched_bias,  # bias
+                    args_list[3],  # activation
+                    args_list[4],  # m
+                    args_list[5],  # n
+                    args_list[6],  # k
+                    args_list[7],  # group_size
+                    bias_stride,  # bias_stride
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_activation(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias
+    as activation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...)
+    lv2 = relax.silu(lv1)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., "silu", ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    pattern = (
+        is_op("relax.nn.silu")(decode_matmul)
+        | is_op("relax.nn.gelu")(decode_matmul)
+        | is_op("relax.nn.relu")(decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int":
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 8
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    matched_activation.op.name[9:],  # activation
+                    args_list[4],  # m
+                    args_list[5],  # n
+                    args_list[6],  # k
+                    args_list[7],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias":
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 10
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    matched_activation.op.name[9:],  # activation
+                    args_list[5],  # m
+                    args_list[6],  # n
+                    args_list[7],  # k
+                    args_list[8],  # group_size
+                    args_list[9],  # bias_stride
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_residual_binary(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.add/multiply` into fastertransformer.gemm_fp16_int_bias as
+    residual binary operation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...)
+    lv2 = relax.add(lv1, residual)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed(
+        "fastertransformer.gemm_fp16_int_bias_residual",
+        ...,
+        residual,
+        ...,
+        "plus",
+        ...
+    )
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    residual = wildcard()
+    pattern = (
+        is_op("relax.add")(decode_matmul, residual)
+        | is_op("relax.add")(residual, decode_matmul)
+        | is_op("relax.multiply")(decode_matmul, residual)
+        | is_op("relax.multiply")(residual, decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias":
+            matched_binary = match[pattern]
+            assert matched_binary.op.name in ["relax.add", "relax.multiply"]
+            binary_op = "plus" if matched_binary.op.name == "relax.add" else "multiply"
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 10
+            matched_residual = match[residual]
+            if not args_list[9].value == 0:
+                # fastertransformer.gemm_fp16_int_bias_residual does not support
+                # bias_stride != 0 yet
+                return expr
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias_residual",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    matched_residual,  # residual
+                    args_list[4],  # activation
+                    binary_op,  # binary_op
+                    "identity",  # unary_op
+                    args_list[5],  # m
+                    args_list[6],  # n
+                    args_list[7],  # k
+                    args_list[8],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_residual_unary(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias_residual
+    as residual unary operation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ...)
+    lv2 = relax.silu(lv1)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ..., "silu", ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    pattern = (
+        is_op("relax.nn.silu")(decode_matmul)
+        | is_op("relax.nn.gelu")(decode_matmul)
+        | is_op("relax.nn.relu")(decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if (
+            match[decode_matmul].args[0].global_symbol
+            == "fastertransformer.gemm_fp16_int_bias_residual"
+        ):
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 12
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias_residual",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    args_list[4],  # residual
+                    args_list[5],  # activation
+                    args_list[6],  # binary_op
+                    matched_activation.op.name[9:],  # activation
+                    args_list[8],  # m
+                    args_list[9],  # n
+                    args_list[10],  # k
+                    args_list[11],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)

python/mlc_llm/compiler_pass/fuse_transpose_matmul.py

@@ -0,0 +1,145 @@
+"""A compiler pass that fuses transpose + matmul."""
+
+import tvm
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
+class FuseTransposeMatmul:
+    """A compiler pass that fuses transpose + matmul."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        mod = relax.transform.FuseOpsByPattern(
+            [
+                (
+                    "transpose_matmul_fuse",
+                    *_pattern(),
+                ),
+            ]
+        )(mod)
+        transpose_matmul_codegen = _TransposeMatmulFuser(mod)
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = transpose_matmul_codegen.visit_expr(func)
+                transpose_matmul_codegen.builder_.update_func(g_var, func)
+        return transpose_matmul_codegen.builder_.get()
+
+
+def _pattern():
+    """Pattern for transpose + matmul."""
+    w = wildcard()
+    x = wildcard()
+    wT = is_op("relax.permute_dims")(w)
+    o = is_op("relax.matmul")(x, wT)
+    annotations = {"o": o, "w": w, "x": x, "wT": wT}
+
+    def _check(context: relax.transform.PatternCheckContext) -> bool:
+        transpose_call = context.annotated_expr["wT"]
+        ndim = transpose_call.args[0].ty.ndim
+        if ndim == -1:
+            return False
+        if ndim == 2 and transpose_call.attrs.axes is None:
+            return True
+        axes = list(range(ndim))
+        axes[-1], axes[-2] = axes[-2], axes[-1]
+        return list(transpose_call.attrs.axes) == axes
+
+    return o, annotations, _check
+
+
+@mutator
+class _TransposeMatmulFuser(PyExprMutator):
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_call_(
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        out_dtype = None
+
+        def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
+            nonlocal out_dtype
+            a_shape = list(a.shape)
+            b_shape = list(b.shape)
+            a_prepended = False
+            b_appended = False
+            if len(a_shape) == 1:
+                a_prepended = True
+                a_shape.insert(0, 1)
+            if len(b_shape) == 1:
+                b_appended = True
+                b_shape.append(1)
+
+            is_a_larger = len(a_shape) > len(b_shape)
+            offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape)
+
+            a_relax = relax.Var("a", relax.TensorType(a.shape))
+            bT_shape = list(b.shape)
+            bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1]
+            bT_relax = relax.Var("b", relax.TensorType(bT_shape))
+            output_shape = self.builder_.normalize(relax.op.matmul(a_relax, bT_relax)).ty.shape
+
+            def matmul_compute(*idx_spatial):
+                k = te.reduce_axis((0, a_shape[-1]), name="k")
+
+                def multiply_compute(idx_reduce):
+                    a_indices = []
+                    b_indices = []
+
+                    for i in range(offset):
+                        if is_a_larger:
+                            a_indices.append(idx_spatial[i])
+                        else:
+                            b_indices.append(idx_spatial[i])
+                    for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)):
+                        a_dim = a_shape[i if is_a_larger else i - offset]
+                        b_dim = b_shape[i if not is_a_larger else i - offset]
+                        dim_equal = a_dim == b_dim
+                        if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0:
+                            a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1
+                            b_dim_is_one = isinstance(b_dim, tirx.IntImm) and b_dim == 1
+                            a_indices.append(0 if a_dim_is_one else idx_spatial[i])
+                            b_indices.append(0 if b_dim_is_one else idx_spatial[i])
+                        else:
+                            a_indices.append(idx_spatial[i])
+                            b_indices.append(idx_spatial[i])
+
+                    if not a_prepended:
+                        a_indices.append(idx_spatial[-2 + b_appended])
+                    a_indices.append(idx_reduce)
+                    if not b_appended:
+                        b_indices.append(idx_spatial[-1])
+                    b_indices.append(idx_reduce)
+
+                    dtype = out_dtype
+                    if dtype != "":
+                        return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
+                    return a(*a_indices) * b(*b_indices)
+
+                return te.sum(multiply_compute(k), axis=k)
+
+            return te.compute(
+                output_shape,
+                lambda *idx: matmul_compute(*idx),
+                name="NT_matmul",
+            )
+
+        if isinstance(call.op, relax.GlobalVar):
+            function = self.builder_.get()[call.op]
+            if (
+                "Composite" in function.attrs
+                and function.attrs["Composite"] == "transpose_matmul_fuse"
+            ):
+                out_dtype = function.ret_ty.dtype
+                return self.builder_.call_te(
+                    te_transposed_matmul,
+                    call.args[1],
+                    call.args[0],
+                    primfunc_name_hint="NT_matmul",
+                )
+
+        return super().visit_call_(call)

python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py

@@ -0,0 +1,198 @@
+"""A compiler pass that lifts TIR-level global allocation to Relax."""
+
+from typing import Dict, List, Tuple  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="LiftTIRGlobalBufferAlloc")
+class LiftTIRGlobalBufferAlloc:
+    """A compiler pass that lifts TIR-level global allocation to Relax."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        return _TIRGlobalAllocRewriter(mod).transform()
+
+
+@mutator
+class _TIRGlobalAllocRewriter(PyExprMutator):
+    def __init__(self, mod: IRModule):
+        super().__init__(mod)
+        self.mod = mod
+        self.gv2new_tensor_sinfo: Dict[  # noqa: UP006
+            tvm.ir.GlobalVar,
+            Tuple[tvm.ir.GlobalVar, List[relax.TensorType], tirx.PrimFunc],  # noqa: UP006
+        ] = {}
+
+    def transform(self) -> IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                updated_func, tensor_sinfo_list = remove_global_buf_alloc(func)
+                if len(tensor_sinfo_list) > 0:
+                    new_gv = self.builder_.add_func(updated_func, g_var.name_hint)
+                    self.gv2new_tensor_sinfo[g_var] = (new_gv, tensor_sinfo_list, func)
+
+        self.mod = self.builder_.get()
+        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)
+
+        mod = self.builder_.get()
+        return relax.transform.DeadCodeElimination()(mod)
+
+    def visit_call_(self, call: relax.Call):
+        call = self.visit_expr_post_order(call)
+        if (
+            call.op != tvm.ir.Op.get("relax.call_tir")
+            or call.args[0] not in self.gv2new_tensor_sinfo
+        ):
+            return call
+
+        g_var = call.args[0]
+        new_gv, tensor_sinfo, func_before_update = self.gv2new_tensor_sinfo[g_var]
+
+        assert len(call.ty_args) == 1
+        if any(_has_symbolic_var(sinfo) for sinfo in tensor_sinfo):
+            tensor_sinfo, success = _resolve_tir_var_mapping(func_before_update, call, tensor_sinfo)
+            if not success:
+                # Cannot resolve TIR var mapping. Fall back to no lifting.
+                self.gv2new_tensor_sinfo.pop(g_var)
+                return call
+
+        args = list(call.args)
+        args[0] = new_gv
+        if isinstance(call.ty_args[0], relax.TensorType):
+            new_call = relax.Call(
+                call.op,
+                args=args,
+                ty_args=[relax.TupleType(list(call.ty_args) + tensor_sinfo)],
+                attrs=call.attrs,
+            )
+            emitted_tuple = self.builder_.emit(new_call)
+            return relax.TupleGetItem(emitted_tuple, 0)
+        assert isinstance(call.ty_args[0], relax.TupleType)
+        return relax.Call(
+            call.op,
+            args=args,
+            ty_args=[relax.TupleType(list(call.ty_args[0].fields) + tensor_sinfo)],
+            attrs=call.attrs,
+        )
+
+
+def remove_global_buf_alloc(
+    func: tirx.PrimFunc,
+) -> Tuple[tirx.PrimFunc, List[relax.TensorType]]:  # noqa: UP006
+    """Remove the global buffer allocation for a given TIR PrimFunc."""
+    assert isinstance(func.body, tirx.SBlockRealize)
+    params = list(func.params)
+    buffer_map = dict(func.buffer_map)
+    tensor_sinfo = []
+    alloc_buffers = []
+
+    insertion_point = len(params)
+    while not isinstance(params[insertion_point - 1].ty, tvm.ir.PointerType):
+        insertion_point -= 1
+        assert insertion_point >= 1
+
+    prev_root_block = func.body.block
+    for buf_alloc in func.body.block.alloc_buffers:
+        if buf_alloc.scope() == "global":
+            param = tirx.Var("var_" + buf_alloc.name, "handle")
+            params.insert(insertion_point, param)
+            insertion_point += 1
+            buffer_map[param] = buf_alloc
+            tensor_sinfo.append(relax.TensorType(buf_alloc.shape, buf_alloc.dtype))
+        else:
+            alloc_buffers.append(buf_alloc)
+
+    if len(tensor_sinfo) == 0:
+        return func, []
+
+    assert len(prev_root_block.iter_vars) == 0
+    assert len(prev_root_block.reads) == 0
+    assert len(prev_root_block.writes) == 0
+    assert len(prev_root_block.match_buffers) == 0
+    assert prev_root_block.name_hint == "root"
+    assert prev_root_block.init is None
+    root_block = tirx.SBlock(
+        iter_vars=[],
+        reads=[],
+        writes=[],
+        name_hint="root",
+        body=prev_root_block.body,
+        alloc_buffers=alloc_buffers,
+        annotations=prev_root_block.annotations,
+    )
+
+    updated_func = tirx.PrimFunc(
+        params=params,
+        body=tirx.SBlockRealize(iter_values=[], predicate=True, block=root_block),
+        ret_type=func.ret_type,
+        buffer_map=buffer_map,
+        attrs=func.attrs,
+    )
+    return updated_func, tensor_sinfo
+
+
+def _has_symbolic_var(tensor_sinfo: relax.TensorType) -> bool:
+    assert isinstance(tensor_sinfo.shape, relax.ShapeExpr)
+    for dim in tensor_sinfo.shape.values:
+        if not isinstance(dim, tirx.IntImm):
+            return True
+    return False
+
+
+def _resolve_tir_var_mapping(
+    func: tirx.PrimFunc,
+    call: relax.Call,
+    tensor_sinfo: List[relax.TensorType],  # noqa: UP006
+) -> Tuple[List[relax.TensorType], bool]:  # noqa: UP006
+    """Resolve the TIR symbolic var relationship across sides of PrimFunc and Relax Function"""
+    var_map: Dict[tirx.Var, tirx.Expr] = {}  # noqa: UP006
+
+    n_arg = len(call.args[1].fields)
+    for i in range(n_arg):
+        buffer_shape = func.buffer_map[func.params[i]].shape
+        arg_shape = call.args[1][i].ty.shape.values
+        assert len(buffer_shape) == len(arg_shape)
+        for v_l, v_r in zip(buffer_shape, arg_shape):
+            if isinstance(v_l, tirx.Var):
+                var_map[v_l] = v_r
+            elif not isinstance(v_l, tirx.IntImm):
+                return [], False
+
+    ret_tensors = call.ty_args[0]
+    ret_tensors = (
+        [ret_tensors] if isinstance(ret_tensors, relax.TensorType) else list(ret_tensors.fields)
+    )
+    for i, ret_tensor in enumerate(ret_tensors):
+        buffer_shape = func.buffer_map[func.params[n_arg + i]].shape
+        ret_tensor_shape = ret_tensor.shape.values
+        assert len(buffer_shape) == len(ret_tensor_shape)
+        for v_l, v_r in zip(buffer_shape, ret_tensor_shape):
+            if isinstance(v_l, tirx.Var):
+                var_map[v_l] = v_r
+            elif not isinstance(v_l, tirx.IntImm):
+                return [], False
+
+    updated_tensor_sinfo = []
+    for sinfo in tensor_sinfo:
+        if not _has_symbolic_var(sinfo):
+            updated_tensor_sinfo.append(sinfo)
+            continue
+        new_shape = []
+        for dim in sinfo.shape.values:
+            new_shape.append(tirx.stmt_functor.substitute(dim, var_map))
+        updated_tensor_sinfo.append(relax.TensorType(new_shape, sinfo.dtype))
+    return updated_tensor_sinfo, True

python/mlc_llm/compiler_pass/low_batch_specialization.py

@@ -0,0 +1,64 @@
+"""A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+import tvm
+import tvm_ffi
+from tvm import tirx
+from tvm.ir.module import IRModule
+from tvm.s_tir import dlight as dl
+
+
+@tvm.transform.module_pass(opt_level=0, name="LowBatchGemvSpecialize")
+class LowBatchGemvSpecialize:
+    """A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                low_batch_range = [2, 8]
+                buckets = [2, 4]
+                low_batch_funcs = []
+                for bucket in buckets:
+                    low_batch_mod = IRModule({})
+                    low_batch_mod["main"] = func
+                    low_batch_mod = dl.ApplyDefaultSchedule(
+                        dl.gpu.LowBatchGEMV(bucket),
+                    )(low_batch_mod)
+                    low_batch_funcs.append(low_batch_mod["main"])
+                if any(
+                    tvm_ffi.structural_equal(low_batch_func, func)
+                    for low_batch_func in low_batch_funcs
+                ):
+                    continue
+                buffers = func.buffer_map.values()
+                shapes = [buffer.shape for buffer in buffers]
+                symbolic_vars = set(
+                    expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var)
+                )
+                if len(symbolic_vars) != 1:
+                    continue
+                gemm_mod = IRModule({})
+                gemm_mod["main"] = func
+                gemm_mod = dl.ApplyDefaultSchedule(
+                    dl.gpu.Matmul(),
+                )(gemm_mod)
+                gemm_func = gemm_mod["main"]
+                sym_var = next(iter(symbolic_vars))
+                body = gemm_func.body
+                for i, range_limit in reversed(list(enumerate(low_batch_range))):
+                    body = tirx.IfThenElse(
+                        tirx.op.tvm_thread_invariant(sym_var <= range_limit),
+                        low_batch_funcs[i].body,
+                        body,
+                    )
+                body = tirx.SBlock([], [], [], "root", body)
+                body = tirx.SBlockRealize([], True, body)
+                new_func = func.with_body(body)
+                new_func = new_func.with_attr("tirx.is_scheduled", 1)
+                new_func = new_func.with_attr("tirx.HoistIfThenElseExprWithBlock", 1)
+                mod.update_func(g_var, new_func)
+        return mod

python/mlc_llm/compiler_pass/pipeline.py

@@ -0,0 +1,209 @@
+"""The compilation pipeline for LLM applications."""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+import tvm
+from tvm import IRModule
+from tvm.relax import register_pipeline
+from tvm.relax.frontend import nn
+from tvm.s_tir import dlight as dl
+
+from mlc_llm.interface.compiler_flags import IPCAllReduceStrategyType
+from mlc_llm.support import logging
+
+from .attach_cuda_graph_alloc_init_func import AttachCUDAGraphAllocInitFunc
+from .attach_embedding_allocator import AttachAllocEmbeddingTensorFunc
+from .attach_logit_processor import AttachLogitProcessFunc
+from .attach_sampler import AttachGPUSamplingFunc
+from .attach_softmax_with_temperature import AttachSoftmaxWithTemperature
+from .attach_spec_decode_aux_funcs import AttachSpecDecodeAuxFuncs
+from .attach_support_info import (
+    AttachAdditionalPrimFuncs,
+    AttachCUDAGraphSymbolicCaptureHints,
+    AttachMemoryPlanAttr,
+    AttachPipelineParallelStages,
+    AttachSequenceLengthPaddingFactor,
+    AttachVariableBounds,
+)
+from .blas_dispatch import BLASDispatch
+from .clean_up_tir_attrs import CleanUpTIRAttrs
+from .dispatch_kv_cache_creation import DispatchKVCacheCreation
+from .dispatch_triton_kernel import DispatchTritonKernel
+from .estimate_memory_usage import AttachMetadataWithMemoryUsage
+from .fuse_add_norm import FuseAddRMSNorm
+from .fuse_dequantize_matmul_ewise import FuseDequantizeMatmulEwise
+from .fuse_dequantize_take import FuseDequantizeTake
+from .fuse_dequantize_transpose import FuseDequantizeTranspose
+from .fuse_ft_dequantize_matmul_epilogue import FuseFTDequantizeEpilogue
+from .fuse_transpose_matmul import FuseTransposeMatmul
+from .lift_global_buffer_alloc import LiftTIRGlobalBufferAlloc
+from .low_batch_specialization import LowBatchGemvSpecialize
+from .pipeline_parallel_rewrite import PipelineParallelRewrite
+from .scatter_tuple_get_item import ScatterTupleGetItem
+
+logger = logging.getLogger(__name__)
+
+
+@tvm.transform.module_pass(opt_level=0, name="_LogProgress")
+class _LogProgress:
+    """A dummy compiler pass that does nothing but logging."""
+
+    def __init__(self, *args):
+        self.args = args
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """A dummy transformation"""
+        logger.info(*self.args)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="DebugDump")
+class _DebugDump:
+    """A dummy compiler pass that does nothing but logging.
+    Only enabled when debug_dump is not None"""
+
+    def __init__(self, file_name: str, file_path: Optional[Path], show_meta: bool = False):
+        self.file_name = file_name
+        self.file_path = file_path
+        self.show_meta = show_meta
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """A dummy transformation that dumps the module to file"""
+        if self.file_path is not None:
+            # NOTE: We use debug level here to avoid spamming the console
+            logger.debug("Dumping IR to %s", self.file_path / self.file_name)
+            with open(self.file_path / self.file_name, "w", encoding="utf-8") as f:
+                f.write(mod.script(show_meta=self.show_meta))
+        return mod
+
+
+@register_pipeline("mlc_llm")
+def _mlc_llm_pipeline(
+    target: tvm.target.Target,
+    flashinfer: bool = False,
+    cublas_gemm: bool = False,
+    faster_transformer: bool = False,
+    allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE,
+    variable_bounds: Optional[Dict[str, int]] = None,  # noqa: UP006
+    cuda_graph_symbolic_capture_hints: Optional[Dict[str, List[str]]] = None,  # noqa: UP006
+    additional_tirs: Optional[Dict[str, tvm.tirx.PrimFunc]] = None,  # noqa: UP006
+    metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ext_mods: Optional[List[nn.ExternModule]] = None,  # noqa: UP006
+    debug_dump: Optional[Path] = None,
+):
+    variable_bounds = variable_bounds or {}
+    cuda_graph_symbolic_capture_hints = cuda_graph_symbolic_capture_hints or {}
+    additional_tirs = additional_tirs or {}
+    metadata = metadata or {}
+    ext_mods = ext_mods or []
+    tensor_parallel_shards = metadata.get("tensor_parallel_shards", 1)
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        seq = tvm.transform.Sequential(
+            [
+                # Phase 0. Add additional information for compilation and remove unused Relax func
+                DispatchKVCacheCreation(target, flashinfer, metadata),
+                AttachSoftmaxWithTemperature(target, metadata),
+                AttachVariableBounds(variable_bounds),
+                AttachCUDAGraphSymbolicCaptureHints(cuda_graph_symbolic_capture_hints),
+                AttachPipelineParallelStages(metadata["pipeline_parallel_stages"]),
+                AttachLogitProcessFunc(target),
+                AttachAdditionalPrimFuncs(additional_tirs),
+                AttachAllocEmbeddingTensorFunc(metadata),
+                AttachGPUSamplingFunc(target, variable_bounds),
+                AttachSpecDecodeAuxFuncs(tensor_parallel_shards),
+                AttachMemoryPlanAttr(),
+                AttachSequenceLengthPaddingFactor(target, metadata),
+                tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)),
+                _DebugDump("debug-phase0.py", debug_dump, show_meta=False),
+                # Phase 1. Passes on high-level operator graph
+                _LogProgress("Running TVM Relax graph-level optimizations"),
+                DispatchTritonKernel(target),
+                FuseFTDequantizeEpilogue(),
+                FuseDequantizeTranspose(),
+                BLASDispatch(target) if cublas_gemm else tvm.transform.Sequential([]),
+                (
+                    FuseAddRMSNorm(target=target)
+                    if target.kind.name != "llvm"
+                    else tvm.transform.Sequential([])
+                ),
+                FuseTransposeMatmul(),
+                _DebugDump("debug-phase1.py", debug_dump, show_meta=False),
+                # Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline
+                _LogProgress("Lowering to TVM TIR kernels"),
+                tvm.relax.backend.DispatchSampling(),
+                tvm.relax.backend.DispatchSortScan(),
+                tvm.relax.transform.LegalizeOps(),
+                tvm.relax.transform.AnnotateTIROpPattern(),
+                tvm.relax.transform.FoldConstant(),
+                tvm.relax.transform.FuseOps(),
+                tvm.relax.transform.FuseTIR(),
+                _DebugDump("debug-phase2.py", debug_dump, show_meta=False),
+                # Phase 3. Passes on TIR
+                _LogProgress("Running TVM TIR-level optimizations"),
+                FuseDequantizeMatmulEwise(),
+                FuseDequantizeTake(),
+                tvm.relax.transform.DeadCodeElimination(),
+                CleanUpTIRAttrs(["op_pattern"]),
+                _DebugDump("debug-phase3.py", debug_dump, show_meta=False),
+                # Phase 4. Low-level Optimizations
+                _LogProgress("Running TVM Dlight low-level optimizations"),
+                LowBatchGemvSpecialize(),
+                (
+                    dl.ApplyDefaultSchedule(
+                        dl.gpu.Matmul(),
+                        dl.gpu.GEMV(),
+                        dl.gpu.Reduction(),
+                        dl.gpu.GeneralReduction(),
+                        dl.gpu.Fallback(),
+                    )
+                    if target.kind.name != "llvm"
+                    else dl.ApplyDefaultSchedule(
+                        dl.cpu.GEMV(),
+                    )
+                ),
+                _DebugDump("debug-phase4.py", debug_dump, show_meta=False),
+                _LogProgress("Lowering to VM bytecode"),
+                (
+                    LiftTIRGlobalBufferAlloc()
+                    if target.kind.name != "llvm"
+                    else tvm.transform.Sequential([])
+                ),
+                (
+                    tvm.tirx.transform.ForceNarrowIndexToInt32()
+                    if target.kind.name != "cuda"
+                    else tvm.transform.Sequential([])
+                ),
+                ScatterTupleGetItem(),
+                PipelineParallelRewrite(),
+                tvm.relax.transform.RewriteDataflowReshape(),
+                tvm.relax.transform.ToNonDataflow(),
+                tvm.relax.transform.RemovePurityChecking(),
+                tvm.relax.transform.CallTIRRewrite(),
+                (
+                    tvm.relax.transform.IPCAllReduceRewrite(allreduce_strategy)
+                    if allreduce_strategy != IPCAllReduceStrategyType.NONE
+                    else tvm.transform.Sequential([])
+                ),
+                tvm.relax.transform.StaticPlanBlockMemory(),
+                AttachMetadataWithMemoryUsage(metadata),
+                _DebugDump("debug-phase5.py", debug_dump, show_meta=False),
+                tvm.relax.transform.RewriteCUDAGraph(),
+                AttachCUDAGraphAllocInitFunc(),
+                tvm.relax.transform.LowerGPUIPCAllocStorage(),
+                tvm.relax.transform.LowerAllocTensor(),
+                tvm.relax.transform.KillAfterLastUse(),
+                tvm.relax.transform.LowerRuntimeBuiltin(),
+                tvm.relax.transform.VMShapeLower(),
+                tvm.relax.transform.AttachGlobalSymbol(),
+                _LogProgress("Compiling external modules"),
+                tvm.relax.transform.AttachExternModules(ext_mods),
+                _LogProgress("Compilation complete! Exporting to disk"),
+            ]
+        )
+        mod = seq(mod)
+        return mod
+
+    return _pipeline

python/mlc_llm/compiler_pass/pipeline_parallel_rewrite.py

@@ -0,0 +1,399 @@
+"""A compiler pass that rewrites IR for pipeline parallelism."""
+
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="PipelineParallelRewrite")
+class PipelineParallelRewrite:
+    """A compiler pass that rewrites IR for pipeline parallelism."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        return _PipelineParallelRewriter(mod.clone()).transform()
+
+
+@mutator
+class _PipelineParallelRewriter(PyExprMutator):
+    def __init__(self, mod: IRModule):
+        super().__init__(mod)
+        self.mod = mod
+        self.old_packed_params_var: relax.Var
+        self.new_main_packed_params_var: relax.Var
+        self.new_stage_func_packed_params: relax.Var
+        self.undefined_shape_vars_remap: Dict[tirx.Var, tirx.Var]  # noqa: UP006
+        self.undefined_param_shape_vars_remap: Dict[tirx.Var, tirx.Var]  # noqa: UP006
+
+    def transform(self) -> IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function) or "pipeline_parallel_stages" not in func.attrs:
+                continue
+            num_stages = int(func.attrs["pipeline_parallel_stages"])
+            if num_stages == 1:
+                continue
+
+            pipeline_stages, stage_send_vars, stage_receive_vars = _extract_pipeline_stages(func)
+            assert len(pipeline_stages) == num_stages, (
+                "Number of pipeline stages mismatches: "
+                f"expecting {num_stages} stages, but {len(pipeline_stages)} are found in the IR."
+            )
+
+            required_func_params = _analyze_required_func_params(pipeline_stages, func.params)
+
+            assert "num_input" in func.attrs
+            num_input = int(func.attrs["num_input"])
+            assert (
+                len(func.params) == num_input + 1
+                and isinstance(func.params[num_input], relax.Var)
+                and func.params[num_input].name_hint == "packed_params"
+            ), 'Only the extra "packed_params" parameter is allowed'
+            self.old_packed_params_var = func.params[num_input]
+            self.new_main_packed_params_var = relax.Var("packed_params", relax.ObjectType())
+            for required_params in required_func_params:
+                for i, param in enumerate(required_params):
+                    if param.same_as(self.old_packed_params_var):
+                        required_params.pop(i)
+                        break
+            func_output = func.body.body
+            assert isinstance(func_output, relax.Var)
+
+            stage_func_gvs = []
+            caller_args_list = []
+            for i in range(num_stages):
+                stage_func_gv, caller_args = self._create_stage_func(
+                    g_var.name_hint + f"_stage{i}",
+                    pipeline_stages[i],
+                    required_func_params[i],
+                    stage_receive_vars[i],
+                    stage_send_vars[i],
+                    func.attrs,
+                    func_output=func_output if i == num_stages - 1 else None,
+                )
+                stage_func_gvs.append(stage_func_gv)
+                caller_args_list.append(caller_args)
+
+            # Create and update the entry function, which dispatches toz the stage functions
+            # according to the disco worker group id.
+            bb = relax.BlockBuilder()
+            params = [*list(func.params[:-1]), self.new_main_packed_params_var]
+            with bb.function(g_var.name_hint, params=params):
+                dispatch_func_args = []
+                for stage_func_gv, caller_args in zip(stage_func_gvs, caller_args_list):
+                    dispatch_func_args.append([stage_func_gv, *caller_args])
+                output = bb.emit(
+                    relax.op.call_builtin_with_ctx(
+                        "mlc.multi_gpu.DispatchFunctionByGroup",
+                        args=[dispatch_func_args],
+                        ty_args=relax.ObjectType(),
+                    )
+                )
+                dispatch_func_gv = bb.emit_func_output(output)
+            dispatch_func = bb.finalize()[dispatch_func_gv]
+            self.builder_.update_func(g_var, dispatch_func)
+
+        return self.builder_.finalize()
+
+    def _create_stage_func(
+        self,
+        func_name: str,
+        stage_bindings: List[relax.Binding],  # noqa: UP006
+        required_func_params: List[relax.Var],  # noqa: UP006
+        stage_receive_vars: List[relax.Var],  # noqa: UP006
+        stage_send_vars: List[relax.Var],  # noqa: UP006
+        func_attrs: tvm.ir.DictAttrs,
+        func_output: Optional[relax.Var],
+    ) -> Tuple[tvm.ir.GlobalVar, List[relax.Expr]]:  # noqa: UP006
+        self.undefined_shape_vars_remap = {}
+        self.undefined_param_shape_vars_remap = {}
+
+        # Prepare the func parameters (except the shape variables and packed params)
+        params, args = self._prepare_stage_func_params_and_args(required_func_params)
+        for new_param, old_param in zip(params, required_func_params):
+            self.set_var_remap(old_param, new_param)
+        # Create new packed params
+        self.new_stage_func_packed_params = relax.Var("packed_params", relax.ObjectType())
+        self.set_var_remap(self.old_packed_params_var, self.new_stage_func_packed_params)
+
+        new_func_outputs = []
+        with self.builder_.function(func_name, pure=False):
+            with self.builder_.dataflow():
+                # Emit the tensors received from last stage.
+                for receive_var in stage_receive_vars:
+                    new_receive_var = self.builder_.emit(
+                        relax.call_dps_packed(
+                            "runtime.disco.recv_from_prev_group",
+                            args=[],
+                            out_ty=self._update_struct_info(receive_var.ty),
+                        ),
+                        name_hint=receive_var.name_hint,
+                    )
+                    self.set_var_remap(receive_var, new_receive_var)
+                # Process the bindings in this stage.
+                for stage_binding in stage_bindings:
+                    if stage_binding.var in stage_send_vars or stage_binding.var.same_as(
+                        func_output
+                    ):
+                        assert isinstance(stage_binding, relax.VarBinding)
+                        new_var = self.builder_.emit_output(
+                            self.visit_expr(stage_binding.value),
+                            name_hint=stage_binding.var.name_hint,
+                        )
+                        self.set_var_remap(stage_binding.var, new_var)
+                        new_func_outputs.append(new_var)
+                    else:
+                        self.visit_binding(stage_binding)
+            # Emit the calls to send tensors to the next stage.
+            for send_var in stage_send_vars:
+                new_send_var = self.get_var_remap(send_var)
+                self.builder_.emit(
+                    relax.Call(
+                        relax.ExternFunc("runtime.disco.send_to_next_group"),
+                        args=[new_send_var],
+                        ty_args=None,
+                    )
+                )
+            # Create the param for the shape variables.
+            shape_var_params = []
+            shape_var_args = []
+            for (
+                shape_var_arg,
+                shape_var_param,
+            ) in self.undefined_shape_vars_remap.items():
+                if shape_var_arg not in self.undefined_param_shape_vars_remap:
+                    shape_var_params.append(shape_var_param)
+                    shape_var_args.append(shape_var_arg)
+            params.append(relax.Var("s", relax.ShapeType(shape_var_params)))
+            args.append(relax.ShapeExpr(shape_var_args))
+            # Add the packed params.
+            params.append(self.new_stage_func_packed_params)
+            args.append(self.new_main_packed_params_var)
+            # Conclude the function.
+            if func_output is not None:
+                assert len(new_func_outputs) == 1
+            new_gv = self.builder_.emit_func_output(
+                (
+                    new_func_outputs[0]
+                    if len(new_func_outputs) == 1
+                    and isinstance(new_func_outputs[0].ty, relax.TupleType)
+                    else new_func_outputs
+                ),
+                params=params,
+            )
+
+        new_func = (
+            self.builder_.get()[new_gv]
+            .with_attrs(func_attrs)
+            .with_attr("num_input", len(params) - 1)
+            .without_attr("global_symbol")
+            .without_attr("pipeline_parallel_stages")
+        )
+        self.builder_.update_func(new_gv, new_func)
+        return new_gv, args
+
+    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
+        if not isinstance(binding.value, relax.TupleGetItem):
+            super().visit_var_binding_(binding)
+            return
+
+        tuple_value = self.visit_expr(binding.value.tuple_value)
+        if not tuple_value.same_as(self.new_stage_func_packed_params):
+            super().visit_var_binding_(binding)
+            return
+
+        assert isinstance(binding.var.ty, relax.TensorType)
+        cur_num_undefined_param_shape_vars = len(self.undefined_param_shape_vars_remap)
+        new_tensor_struct_info = self._update_struct_info(
+            binding.var.ty, self.undefined_param_shape_vars_remap
+        )
+        has_new_undefined_shape_var = (
+            len(self.undefined_param_shape_vars_remap) != cur_num_undefined_param_shape_vars
+        )
+        self.undefined_shape_vars_remap = {
+            **self.undefined_shape_vars_remap,
+            **self.undefined_param_shape_vars_remap,
+        }
+        ret_sinfo = (
+            new_tensor_struct_info if not has_new_undefined_shape_var else relax.ObjectType()
+        )
+        call = relax.call_pure_packed(
+            "vm.builtin.tuple_getitem",
+            self.new_stage_func_packed_params,
+            relax.prim_value(binding.value.index),
+            ty_args=ret_sinfo,
+        )
+        new_binding_var = self.builder_.emit(call, binding.var.name_hint)
+        if has_new_undefined_shape_var:
+            new_binding_var = self.builder_.match_cast(
+                new_binding_var, new_tensor_struct_info, binding.var.name_hint + "_cast"
+            )
+        self.set_var_remap(binding.var, new_binding_var)
+
+    def visit_call_(self, call: relax.Call) -> relax.Call:
+        call = super().visit_call_(call)
+        return relax.Call(
+            call.op,
+            call.args,
+            call.attrs,
+            ty_args=[self._update_struct_info(struct_info) for struct_info in call.ty_args],
+        )
+
+    def _prepare_stage_func_params_and_args(
+        self,
+        required_func_params: List[relax.Var],  # noqa: UP006
+    ) -> Tuple[List[relax.Var], List[relax.Expr]]:  # noqa: UP006
+        params: List[relax.Var] = []  # noqa: UP006
+        args: List[relax.Expr] = []  # noqa: UP006
+        for required_param in required_func_params:
+            struct_info = self._update_struct_info(required_param.ty)
+            params.append(relax.Var(required_param.name_hint, struct_info))
+            args.append(required_param)
+
+        return params, args
+
+    def _update_struct_info(
+        self,
+        struct_info: relax.Type,
+        undefined_var_remap: Optional[Dict[tirx.Var, tirx.Var]] = None,  # noqa: UP006
+    ) -> relax.Type:
+        if undefined_var_remap is None:
+            undefined_var_remap = self.undefined_shape_vars_remap
+        if isinstance(struct_info, relax.TensorType):
+            return (
+                relax.TensorType(
+                    self._update_shape(struct_info.shape.values, undefined_var_remap),
+                    struct_info.dtype,
+                )
+                if struct_info.shape is not None and isinstance(struct_info.shape, relax.ShapeExpr)
+                else struct_info
+            )
+        if isinstance(struct_info, relax.ShapeType):
+            return (
+                relax.ShapeType(self._update_shape(struct_info.values, undefined_var_remap))
+                if struct_info.values is not None
+                else struct_info
+            )
+        if isinstance(struct_info, relax.ObjectType):
+            return relax.ObjectType()
+        if isinstance(struct_info, relax.TupleType):
+            return relax.TupleType(
+                [self._update_struct_info(field_sinfo) for field_sinfo in struct_info.fields]
+            )
+        return struct_info
+
+    def _copy_undefined_var(
+        self,
+        expr: tirx.Expr,
+        undefined_var_remap: Dict[tirx.Var, tirx.Var],  # noqa: UP006
+    ) -> None:
+        def _visit_expr(e: tirx.Expr) -> None:
+            if isinstance(e, tirx.Var) and e not in undefined_var_remap:
+                new_var = tirx.Var(e.name, e.ty)
+                undefined_var_remap[e] = new_var
+
+        tirx.stmt_functor.post_order_visit(expr, _visit_expr)
+
+    def _update_shape(
+        self,
+        shape: List[tirx.Expr],  # noqa: UP006
+        undefined_var_remap: Dict[tirx.Var, tirx.Var],  # noqa: UP006
+    ) -> List[tirx.Expr]:  # noqa: UP006
+        new_shape = []
+        for v in shape:
+            self._copy_undefined_var(v, undefined_var_remap)
+            new_shape.append(tirx.stmt_functor.substitute(v, undefined_var_remap))
+        return new_shape
+
+
+def _extract_pipeline_stages(
+    func: relax.Function,
+) -> Tuple[List[List[relax.Binding]], List[List[relax.Var]], List[List[relax.Var]]]:  # noqa: UP006
+    pipeline_stages: List[List[relax.Binding]] = []  # noqa: UP006
+    stage_send_vars: List[List[relax.Var]] = []  # noqa: UP006
+    stage_receive_vars: List[List[relax.Var]] = []  # noqa: UP006
+
+    # Requiring that the function has only one body block which is a dataflow block
+    assert isinstance(func.body, relax.SeqExpr)
+    assert len(func.body.blocks) == 1
+    assert isinstance(func.body.blocks[0], relax.DataflowBlock)
+    bindings = func.body.blocks[0].bindings
+
+    boundary_var = None
+    current_stage_bindings: List[relax.Binding] = []  # noqa: UP006
+    current_stage_receive_vars: List[relax.Var] = []  # noqa: UP006
+    for binding in bindings:
+        if (
+            isinstance(binding, relax.VarBinding)
+            and isinstance(binding.value, relax.Call)
+            and binding.value.op == tvm.ir.Op.get("relax.call_pure_packed")
+            and binding.value.args[0].global_symbol == "mlc.pipeline_parallel_stage_boundary"
+        ):
+            assert len(current_stage_bindings) > 0
+            pipeline_stages.append(current_stage_bindings)
+            assert all(receive_var is not None for receive_var in current_stage_receive_vars)
+            stage_receive_vars.append(current_stage_receive_vars)
+            args = binding.value.args[1:]
+            assert len(args) >= 1 and all(isinstance(arg, relax.Var) for arg in args)
+            stage_send_vars.append(list(args))
+
+            boundary_var = binding.var
+            current_stage_bindings = []
+            current_stage_receive_vars = [boundary_var] if len(args) == 1 else [None for _ in args]
+        elif (
+            isinstance(binding, relax.VarBinding)
+            and isinstance(binding.value, relax.TupleGetItem)
+            and binding.value.tuple_value.same_as(boundary_var)
+        ):
+            current_stage_receive_vars[binding.value.index] = binding.var
+        else:
+            current_stage_bindings.append(binding)
+
+    assert len(current_stage_bindings) > 0
+    pipeline_stages.append(current_stage_bindings)
+    assert all(receive_var is not None for receive_var in current_stage_receive_vars)
+    stage_receive_vars.append(current_stage_receive_vars)
+    stage_send_vars.append([])
+
+    return pipeline_stages, stage_send_vars, stage_receive_vars
+
+
+def _analyze_required_func_params(
+    pipeline_stages: List[List[relax.Binding]],  # noqa: UP006
+    func_params: List[relax.Var],  # noqa: UP006
+) -> List[List[relax.Var]]:  # noqa: UP006
+    analyzer = _RequiredFuncParamAnalyzer(func_params)
+    required_func_params: List[List[relax.Var]] = []  # noqa: UP006
+    for stage_bindings in pipeline_stages:
+        required_params: List[relax.Var]  # noqa: UP006
+        required_params = analyzer.run(stage_bindings)
+        required_func_params.append(required_params)
+    return required_func_params
+
+
+@visitor
+class _RequiredFuncParamAnalyzer(PyExprVisitor):
+    """The IR visitor which analyzes the required func parameters in each pipeline stage."""
+
+    def __init__(self, func_params: List[relax.Var]) -> None:  # noqa: UP006
+        self.func_params = set(func_params)
+        self.required_params: List[relax.Var]  # noqa: UP006
+
+    def run(self, stage_bindings: List[relax.Binding]) -> List[relax.Var]:  # noqa: UP006
+        """Entry point of the visitor."""
+        self.required_params = []
+        for binding in stage_bindings:
+            self.visit_binding(binding)
+        return self.required_params
+
+    def visit_var_(self, var: relax.Var) -> None:
+        if var in self.func_params:
+            if var not in self.required_params:
+                self.required_params.append(var)

python/mlc_llm/compiler_pass/scatter_tuple_get_item.py

@@ -0,0 +1,49 @@
+"""A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
+
+from typing import Dict  # noqa: UP035
+
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr import Expr, Var
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="ScatterTupleGetItem")
+class ScatterTupleGetItem:
+    """A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _Scatter(mod).transform()
+
+
+@mutator
+class _Scatter(PyExprMutator):
+    def __init__(self, mod: IRModule) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.var_map: Dict[Var, Expr] = {}  # noqa: UP006
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        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()
+
+    def visit_var_binding_(self, binding: relax.VarBinding):
+        super().visit_var_binding_(binding)
+        if isinstance(binding.value, relax.TupleGetItem):
+            self.var_map[binding.var] = binding.value
+
+    def visit_dataflow_var_(self, var: relax.DataflowVar) -> Expr:
+        if var in self.var_map:
+            new_var = self.builder_.emit(self.var_map[var], name_hint=var.name_hint)
+            self.set_var_remap(var, new_var)
+            self.var_map.pop(var)
+            return new_var
+        return var