chore: import upstream snapshot with attribution

python/mlc_llm/op/__init__.py

@@ -0,0 +1,15 @@
+"""Extern module for compiler."""
+
+from . import moe_matmul, moe_misc
+from .attention import attention
+from .batch_spec_verify import batch_spec_verify
+from .extern import configure, enable, get_store
+from .ft_gemm import faster_transformer_dequantize_gemm
+from .mrope import (
+    MultimodalRotaryEmbedding,
+    VisionPositionMetadata,
+    apply_multimodal_rotary_pos_emb,
+    get_mrope_position_ids,
+)
+from .pipeline_parallel import pipeline_stage_boundary
+from .top_p_pivot import top_p_pivot, top_p_renorm

python/mlc_llm/op/attention.py

@@ -0,0 +1,186 @@
+"""Operators enabled by external modules."""
+
+from typing import Optional
+
+import tvm
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Tensor, op
+
+from mlc_llm.support import logging
+
+from . import extern as _extern
+
+logger = logging.getLogger(__name__)
+
+
+WARN_FLASHINFER_GROUP_SIZE = False
+WARN_FLASHINFER_HEAD_DIM = False
+
+
+def attention(
+    q: nn.Tensor,
+    k: nn.Tensor,
+    v: nn.Tensor,
+    casual_mask: nn.Tensor,
+    attn_score_scaling_factor: float = 1.0,
+    qk_dtype: Optional[str] = None,
+) -> nn.Tensor:
+    """Attention with casual mask.
+
+    --- Variables ---
+    s: sequence length of the current query
+    t: total sequence length
+    d: head dimension
+    h, h_q: number of heads in query
+    h_kv: number of heads in key and value
+    b: batch size = 1
+
+    --- Shapes ---
+    q: [b, s, h_q, d]
+    k: [t, h_kv, d]
+    v: [t, h_kv, d]
+    o: [1, s, hidden = h_q * d]
+
+    --- Computation ---
+
+    .. code-block:: python
+
+        if h_kv != h_q:
+            k = k.repeat(h_q // h_kv, axis=1)
+            v = v.repeat(h_q // h_kv, axis=1)
+        q -> [b, h, s, d]
+        k, v -> [b, h, t, d]
+        attn = q @ k^T / sqrt(d) * attn_score_scaling_factor  # [b, h, s, t]
+        attn = softmax_with_mask(attn, casual_mask, axis=-1)
+        o = attn @ v  # [b, h, s, d]
+        o -> [b, s, h * d]
+
+    --- Other params ---
+    qk_dtype: if set, `matmul(Q, K, out_dtype=qk_dtype)`, (otherwise use `q.dtype` as `out_dtype`).
+        For FlashInfer, if "float32", sets `allow_fp16_qk_reduction` to False; otherwise no effect.
+    """
+    assert q.ndim == 4 and k.ndim in [3, 4] and v.ndim in [3, 4]
+    b, s, h_q, d = q.shape
+    t, h_kv, _ = k.shape[-3:]
+    group_size = h_q // h_kv
+
+    def _fallback():
+        from tvm.relax.frontend.nn.llm.kv_cache import (
+            _attention_sequence_prefill,
+        )
+
+        nonlocal q, k, v, qk_dtype
+        if k.ndim == 3:
+            k = op.reshape(k, [b, t, h_kv, d])
+        if v.ndim == 3:
+            v = op.reshape(v, [b, t, h_kv, d])
+        if h_kv != h_q:
+            k = k.repeat(h_q // h_kv, axis=2)
+            v = v.repeat(h_q // h_kv, axis=2)
+
+        target = tvm.target.Target("cuda")
+        attn_output, _ = op.tensor_ir_op(
+            _attention_sequence_prefill(
+                h_kv=h_kv,
+                h_q=h_q,
+                d=d,
+                dtype=q.dtype,
+                target=target,
+                sm_scale=attn_score_scaling_factor / (d**0.5),
+            ),
+            "sequence_prefill",
+            [q, k, v],
+            [
+                Tensor.placeholder([b, s, h_q, d], q.dtype),
+                Tensor.placeholder([b, s, h_q], q.dtype),
+            ],
+        )
+
+        output = op.reshape(attn_output, shape=(b, s, h_q * d))
+        return output
+
+    # FlashInfer Implementation
+    if (
+        _extern.get_store().flashinfer
+        and attn_score_scaling_factor == 1.0
+        and q.dtype == "float16"
+        and k.dtype == "float16"
+        and v.dtype == "float16"
+    ):
+        if group_size not in [1, 4, 6, 8]:
+            global WARN_FLASHINFER_GROUP_SIZE
+            if not WARN_FLASHINFER_GROUP_SIZE:
+                WARN_FLASHINFER_GROUP_SIZE = True
+                logger.warning(
+                    "FlashInfer only supports group size in [1, 4, 6, 8], but got %d. Skip and "
+                    "fallback to default implementation.",
+                    group_size,
+                )
+            return _fallback()
+        if d not in [128]:
+            global WARN_FLASHINFER_HEAD_DIM
+            if not WARN_FLASHINFER_HEAD_DIM:
+                WARN_FLASHINFER_HEAD_DIM = True
+                logger.warning(
+                    "FlashInfer only supports head_dim in [128], but got %d. Skip and fallback to "
+                    "default implementation.",
+                    d,
+                )
+            return _fallback()
+        rope_theta = 0.0
+        rope_scale = 1.0
+        qkv_layout = 0  # "NHD", N for seq_len, H for num_heads, D for head_dim
+        rotary_mode = 0  # "kNone"
+        casual = 1  # True
+        fp16_qk = 1  # True
+        if qk_dtype == "float32":
+            fp16_qk = 0  # False
+
+        # 32MB scratchpad
+        scratch = op.empty([8192 * 1024], dtype="float32")
+
+        def _decode():
+            return op.extern(
+                name="flashinfer.single_decode",
+                args=[
+                    q,
+                    k,
+                    v,
+                    scratch,
+                    qkv_layout,
+                    rotary_mode,
+                    rope_scale,
+                    rope_theta,
+                ],
+                out=nn.Tensor.placeholder((b, s, h_q * d), dtype="float16"),
+            )
+
+        def _prefill():
+            return op.extern(
+                name="flashinfer.single_prefill",
+                args=[
+                    q,
+                    k,
+                    v,
+                    scratch,
+                    casual,
+                    qkv_layout,
+                    rotary_mode,
+                    fp16_qk,
+                    rope_scale,
+                    rope_theta,
+                ],
+                out=nn.Tensor.placeholder((b, s, h_q * d), dtype="float16"),
+            )
+
+        if isinstance(s, int) and s == 1:
+            func = "decode"
+        else:
+            func = "prefill"
+        return {
+            "decode": _decode,
+            "prefill": _prefill,
+        }[func]()
+
+    # Fallback Implementation
+    return _fallback()

python/mlc_llm/op/batch_matmul.py

@@ -0,0 +1,44 @@
+"""Batch matmul operators"""
+
+from typing import Tuple  # noqa: UP035
+
+from tvm.relax.frontend import nn
+
+from mlc_llm.op import cutlass
+from mlc_llm.quantization.block_scale_quantization import rowwise_group_quant_fp8
+
+
+def quantized_bmm(
+    x: nn.Tensor,
+    w: nn.Tensor,
+    w_scale: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+) -> nn.Tensor:
+    """Quantized batch matmul.
+    Currently only support CUDA backend (by using CUTLASS).
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [b, m, k].
+
+    w : nn.Tensor
+        The weight tensor, with shape of [b, n, k] (column major).
+
+    w_scale : nn.Tensor
+        The scale tensor, with shape of [b, n // block_size[0], k // block_size[1]].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    Returns
+    -------
+    ret : nn.Tensor
+        The output tensor, with shape of [b, m, n].
+    """
+    x_fp8, x_scale = rowwise_group_quant_fp8(
+        x, block_size[1], w.dtype, transpose_scale=True, keep_first_batch_dim=True
+    )
+    return cutlass.fp8_groupwise_scaled_bmm(
+        x_fp8, x_scale, w, w_scale, block_size, out_dtype=x.dtype
+    )

python/mlc_llm/op/batch_spec_verify.py

@@ -0,0 +1,175 @@
+"""Operators for batch verify in speculative decoding."""
+
+from tvm.script import tirx as T
+
+# mypy: disable-error-code="attr-defined,valid-type,name-defined"
+
+
+def batch_spec_verify(vocab_size):
+    """Batch draft verify function. This function verifies the token tree.
+
+    Before calling the function
+
+    - token_tree_parent_ptr[b] should store the root of the tree
+
+    - draft_probs[node_id, :] stores the prob that samples the correspond tree node
+    - model_probs[node_id, :] stores the prob that should be used to sample its children
+    - Please note that the storage convention difference between model_probs and draft_probs
+        draft_probs was stored on the token node, while model_probs stores on the parent.
+        This is an intentional design since we can sample different child token with different
+        proposal draft probabilities, but the ground truth model_prob is unique per parent.
+
+    After calling the function
+    - token_tree_parent_ptr[b] points to the last token accepted
+    - There should be a followup sample step that samples from model_probs[token_tree_parent_ptr[b], :]
+        This token will be appended to the token generated.
+
+    This function will inplace update model_probs if a token was rejected and renormalization is needed.
+
+    Parameters
+    ----------
+    draft_probs:
+        The draft probability attached to each tree node
+
+    draft_tokens:
+        The draft token in each node
+
+    model_probs:
+        The model proability attached to each parent
+
+    token_tree_first_child:
+        The first child of each tree node, if there is no child, it should be -1
+
+    token_tree_next_sibling
+        The next sibling of each tree node, if there is no next sibling, it should be -1
+
+    uniform_samples
+        Per node uniform sample used to check rejection
+
+    token_tree_parent_ptr:
+        Current parent ptr state
+    """  # noqa: E501
+    TX = 1024
+
+    def _var(dtype="int32"):
+        return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+    # fmt: off
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_draft_probs: T.handle,
+        var_draft_tokens: T.handle,
+        var_model_probs: T.handle,
+        var_token_tree_first_child: T.handle,
+        var_token_tree_next_sibling: T.handle,
+        var_uniform_samples: T.handle,
+        var_token_tree_parent_ptr: T.handle,
+    ):
+        """
+        [
+            blockIdx.x on batch,
+            threadIdx.x on vocab_size,
+            for loop over excessive amounts
+        ]
+        """
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        num_nodes = T.int32()
+        nbatch = T.int32()
+
+        draft_probs = T.match_buffer(var_draft_probs, (num_nodes, vocab_size), "float32")
+        draft_tokens = T.match_buffer(var_draft_tokens, (num_nodes,), "int32")
+        model_probs = T.match_buffer(var_model_probs, (num_nodes, vocab_size), "float32")
+        token_tree_first_child = T.match_buffer(var_token_tree_first_child, (num_nodes,), "int32")
+        token_tree_next_sibling = T.match_buffer(var_token_tree_next_sibling, (num_nodes,), "int32")
+        uniform_samples = T.match_buffer(var_uniform_samples, (num_nodes,), "float32")
+        token_tree_parent_ptr = T.match_buffer(var_token_tree_parent_ptr, (nbatch,), "int32")
+
+        with T.sblock("kernel"):
+            child_ptr = _var()
+            parent_ptr = _var()
+            child_token = _var()
+            done = _var("bool")
+            psum = _var("float32")
+            t0 = _var("float32")
+            model_prob_local = _var("float32")
+            draft_prob_local = _var("float32")
+            p_child = _var("float32")
+            q_child = _var("float32")
+            uniform_sample = _var("float32")
+
+            pred_shared = T.sblock_alloc_buffer((1,), "bool", scope="shared")
+            pred_local = T.sblock_alloc_buffer((1,), "bool", scope="local")
+
+            for _bx in T.thread_binding(0, nbatch, thread="blockIdx.x"):
+                for _tx in T.thread_binding(0, TX, thread="threadIdx.x"):
+                    with T.sblock("CTA"):
+                        # batch size
+                        b = T.axis.S(nbatch, _bx)
+                        tx = T.axis.S(TX, _tx)
+
+                        parent_ptr[0] = token_tree_parent_ptr[b]
+                        child_ptr[0] = token_tree_first_child[parent_ptr[0]]
+                        done[0] = False
+
+                        while T.Not(done[0]):
+                            T.tvm_storage_sync("shared") # ensure all effects last round are visible
+                            if child_ptr[0] == -1:
+                                done[0] = True
+                                T.tvm_storage_sync("shared") # sync before exit
+                            else:
+                                # decide to validate current ptr
+                                if tx == 0:
+                                    child_token[0] = draft_tokens[child_ptr[0]]
+                                    p_child[0] = model_probs[parent_ptr[0], child_token[0]]
+                                    q_child[0] = draft_probs[child_ptr[0], child_token[0]]
+                                    uniform_sample[0] = uniform_samples[child_ptr[0]]
+                                    pred_shared[0] = p_child[0] >= uniform_sample[0] * q_child[0]  # use multiplication to avoid division by zero  # noqa: E501
+                                T.tvm_storage_sync("shared") # make sure all read of model_probs are done  # noqa: E501
+                                pred_local[0] = pred_shared[0]
+
+                                # accept the proposal, we move to child
+                                if pred_local[0]:
+                                    parent_ptr[0] = child_ptr[0]
+                                    child_ptr[0] = token_tree_first_child[child_ptr[0]]
+                                else:
+                                    psum[0] = 0.0
+                                    # renormalize probability, predicated by stopped_expansion[b]:
+                                    for i in T.serial(T.ceildiv(vocab_size, TX)):
+                                        k = T.meta_var(i * TX + tx)
+                                        if k < vocab_size:
+                                            model_prob_local[0] = model_probs[parent_ptr[0], k]
+                                            draft_prob_local[0] = draft_probs[child_ptr[0], k]
+                                            model_prob_local[0] = T.max(model_prob_local[0] - draft_prob_local[0], 0.0)  # noqa: E501
+                                            psum[0] += model_prob_local[0]
+
+                                    with T.sblock("block_cross_thread"):
+                                        T.reads(psum[0])
+                                        T.writes(t0[0])
+                                        T.attr(
+                                            T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                                            "reduce_scope",
+                                            T.int32(0),
+                                        )
+                                        T.tvm_thread_allreduce(T.uint32(1), psum[0], True, t0[0], tx, dtype="void")  # noqa: E501
+
+                                    if t0[0] < 1e-7:
+                                        # accept the proposal, we move to child
+                                        parent_ptr[0] = child_ptr[0]
+                                        child_ptr[0] = token_tree_first_child[child_ptr[0]]
+                                    else:
+                                        # renormalize
+                                        for i in T.serial(T.ceildiv(vocab_size, TX)):
+                                            k = T.meta_var(i * TX + tx)
+                                            if k < vocab_size:
+                                                model_prob_local[0] = model_probs[parent_ptr[0], k]
+                                                draft_prob_local[0] = draft_probs[child_ptr[0], k]
+                                                model_prob_local[0] = T.max(model_prob_local[0] - draft_prob_local[0], 0.0)  # noqa: E501
+                                                model_probs[parent_ptr[0], k] = model_prob_local[0] / t0[0]  # noqa: E501
+
+                                        child_ptr[0] = token_tree_next_sibling[child_ptr[0]]
+
+                        if tx == 0:
+                            token_tree_parent_ptr[b] = parent_ptr[0]
+    # fmt: on
+
+    return _func

python/mlc_llm/op/cutlass.py

@@ -0,0 +1,375 @@
+"""Operators enabled by external modules."""
+
+from typing import Optional, Tuple  # noqa: UP035
+
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import op
+
+
+def group_gemm(
+    x: nn.Tensor,
+    weight: nn.Tensor,
+    indptr: nn.Tensor,
+    scale: Optional[nn.Tensor] = None,
+    weight_dtype: Optional[str] = None,
+    out_dtype: Optional[str] = None,
+):
+    """
+    Cutlass group gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [num_groups, n, k].
+
+    indptr : nn.Tensor
+        The indptr tensor, with shape of [num_groups].
+
+    scale : Optional[nn.Tensor]
+        The scale tensor, with shape of [1].
+
+    weight_dtype: Optional[str]
+        The data type of the weight tensor.
+
+    out_dtype: Optional[str]
+        The data type of the output tensor.
+
+    Returns
+    -------
+    nn.Tensor
+        The output tensor, with shape of [m, n].
+    """
+    assert x.ndim == 2
+    assert weight.ndim == 3
+    assert indptr.ndim == 1
+    assert weight.shape[0] == indptr.shape[0]
+    assert indptr.dtype == "int64"
+    out_dtype = out_dtype if out_dtype else x.dtype
+    weight_dtype = weight_dtype if weight_dtype else weight.dtype
+
+    if x.dtype == "float8_e5m2" and weight_dtype == "float8_e5m2" and out_dtype == "float16":
+        func_name = "cutlass.group_gemm_e5m2_e5m2_fp16"
+    elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e5m2" and out_dtype == "float16":
+        func_name = "cutlass.group_gemm_e4m3_e5m2_fp16"
+    elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e4m3fn" and out_dtype == "float16":
+        func_name = "cutlass.group_gemm_e4m3_e4m3_fp16"
+    elif (x.dtype == "float16" and weight_dtype == "float16" and out_dtype == "float16") or (
+        x.dtype == "bfloat16" and weight_dtype == "bfloat16" and out_dtype == "bfloat16"
+    ):
+        func_name = "cutlass.group_gemm"
+    else:
+        raise NotImplementedError(
+            f"Unsupported data type: x={x.dtype}, weight={weight_dtype}, out={out_dtype}"
+        )
+
+    if "float8" in x.dtype:
+        assert scale is not None, "scale is required for float8 input"
+
+    workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace")
+
+    return op.extern(
+        func_name,
+        args=[x, weight, indptr, workspace] + ([scale] if scale is not None else []),
+        out=nn.Tensor.placeholder((x.shape[0], weight.shape[1]), dtype=out_dtype),
+    )
+
+
+def fp8_gemm(
+    x: nn.Tensor,
+    weight: nn.Tensor,
+    scale: nn.Tensor,
+    weight_dtype: Optional[str] = None,
+    out_dtype: Optional[str] = None,
+):
+    """
+    Cutlass fp8 gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [num_groups, n, k].
+
+    scale : Optional[nn.Tensor]
+        The scale tensor, with shape of [1].
+
+    weight_dtype: Optional[str]
+        The data type of the weight tensor.
+
+    out_dtype: Optional[str]
+        The data type of the output tensor.
+
+    Returns
+    -------
+    nn.Tensor
+        The output tensor, with shape of [m, n].
+    """
+    assert x.ndim >= 2
+    assert weight.ndim == 2
+    assert scale.ndim == 1 and scale.shape[0] == 1
+    out_dtype = out_dtype if out_dtype else x.dtype
+    weight_dtype = weight_dtype if weight_dtype else weight.dtype
+
+    if x.dtype == "float8_e5m2" and weight_dtype == "float8_e5m2" and out_dtype == "float16":
+        func_name = "cutlass.gemm_e5m2_e5m2_fp16"
+    elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e5m2" and out_dtype == "float16":
+        func_name = "cutlass.gemm_e5m2_e4m3_fp16"
+    elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e4m3fn" and out_dtype == "float16":
+        func_name = "cutlass.gemm_e4m3_e4m3_fp16"
+    else:
+        raise NotImplementedError(
+            f"Unsupported data type: x={x.dtype}, weight={weight_dtype}, out={out_dtype}"
+        )
+
+    workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace")
+
+    return op.extern(
+        func_name,
+        args=[x, weight, workspace, scale],
+        out=nn.Tensor.placeholder((*x.shape[:-1], weight.shape[0]), dtype=out_dtype),
+    )
+
+
+def fp8_groupwise_scaled_gemm(
+    x: nn.Tensor,
+    x_scale: nn.Tensor,
+    weight: nn.Tensor,
+    weight_scale: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+):
+    """Cutlass block-scale fp8 gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    x_scale : nn.Tensor
+        The scale tensor, with shape of [k // block_size, m].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [n, k].
+
+    weight_scale : nn.Tensor
+        The scale tensor, with shape of [n // block_size, k // block_size].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    out_dtype : str
+        The data type of the output tensor.
+
+    Returns
+    -------
+    out : nn.Tensor
+        The output tensor, with shape of [m, n] and dtype of `out_dtype`.
+    """
+    assert x.ndim >= 2
+    assert weight.ndim == 2
+    assert x_scale.ndim == x.ndim
+    assert weight_scale.ndim == weight.ndim
+
+    if block_size[0] != 128 or block_size[1] != 128:
+        raise ValueError(f"block_size must be (128, 128), but got {block_size}")
+    if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn":
+        raise ValueError(
+            f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}"
+        )
+    if x_scale.dtype != "float32" or weight_scale.dtype != "float32":
+        raise ValueError(
+            "x_scale and weight_scale must be float32, but got "
+            f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}"
+        )
+    if out_dtype not in ["float16", "bfloat16"]:
+        raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}")
+
+    func_name = "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn"
+    workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace")
+    return op.extern(
+        func_name,
+        args=[
+            x,
+            weight,
+            x_scale,
+            weight_scale,
+            workspace,
+            block_size[0],
+            block_size[1],
+        ],
+        out=nn.Tensor.placeholder((*x.shape[:-1], weight.shape[0]), dtype=out_dtype),
+    )
+
+
+def fp8_groupwise_scaled_bmm(
+    x: nn.Tensor,
+    x_scale: nn.Tensor,
+    weight: nn.Tensor,
+    weight_scale: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+):
+    """Cutlass block-scale fp8 gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [b, m, k].
+
+    x_scale : nn.Tensor
+        The scale tensor, with shape of [b, k // block_size, m].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [b, n, k].
+
+    weight_scale : nn.Tensor
+        The scale tensor, with shape of [b, n // block_size, k // block_size].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    out_dtype : str
+        The data type of the output tensor.
+
+    Returns
+    -------
+    out : nn.Tensor
+        The output tensor, with shape of [m, n] and dtype of `out_dtype`.
+    """
+    assert x.ndim == 3
+    assert weight.ndim == 3
+    assert x_scale.ndim == x.ndim
+    assert weight_scale.ndim == weight.ndim
+    assert x.shape[0] == x_scale.shape[0] == weight.shape[0] == weight_scale.shape[0]
+
+    if block_size[0] != 128 or block_size[1] != 128:
+        raise ValueError(f"block_size must be (128, 128), but got {block_size}")
+    if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn":
+        raise ValueError(
+            f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}"
+        )
+    if x_scale.dtype != "float32" or weight_scale.dtype != "float32":
+        raise ValueError(
+            "x_scale and weight_scale must be float32, but got "
+            f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}"
+        )
+    if out_dtype not in ["float16", "bfloat16"]:
+        raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}")
+
+    func_name = "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn"
+    workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace")
+    return op.extern(
+        func_name,
+        args=[
+            x,
+            weight,
+            x_scale,
+            weight_scale,
+            workspace,
+            block_size[0],
+            block_size[1],
+        ],
+        out=nn.Tensor.placeholder((x.shape[0], x.shape[1], weight.shape[1]), dtype=out_dtype),
+    )
+
+
+def fp8_groupwise_scaled_group_gemm(
+    x: nn.Tensor,
+    x_scale: nn.Tensor,
+    weight: nn.Tensor,
+    weight_scale: nn.Tensor,
+    indptr: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+):
+    """Triton block-scale fp8 group gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    x_scale : nn.Tensor
+        The scale tensor, with shape of [m, k // block_size].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [num_experts, n, k].
+
+    weight_scale : nn.Tensor
+        The scale tensor, with shape of [num_experts, n // block_size, k // block_size].
+
+    indptr : nn.Tensor
+        The indptr tensor of group gemm, with shape of [num_experts + 1,].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    out_dtype : str
+        The data type of the output tensor.
+
+    Returns
+    -------
+    out : nn.Tensor
+        The output tensor, with shape of [m, n] and dtype of `out_dtype`.
+    """
+    assert x.ndim >= 2
+    assert weight.ndim == 3
+    assert x_scale.ndim == x.ndim
+    assert weight_scale.ndim == weight.ndim
+    assert x.shape[-1] == weight.shape[2]
+    assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1]
+    assert (weight.shape[2] + block_size[1] - 1) // block_size[1] == weight_scale.shape[2]
+    assert (weight.shape[1] + block_size[0] - 1) // block_size[0] == weight_scale.shape[1]
+
+    if block_size[0] != 128 or block_size[1] != 128:
+        raise ValueError(f"block_size must be (128, 128), but got {block_size}")
+    if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn":
+        raise ValueError(
+            f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}"
+        )
+    if x_scale.dtype != "float32" or weight_scale.dtype != "float32":
+        raise ValueError(
+            "x_scale and weight_scale must be float32, but got "
+            f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}"
+        )
+    if out_dtype not in ["float16", "bfloat16"]:
+        raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}")
+
+    num_experts = weight.shape[0]
+    m = x.shape[0]
+    for i in range(1, x.ndim - 1):
+        m *= x.shape[i]
+    n = weight.shape[1]
+    k = x.shape[-1]
+    assert weight_scale.shape[0] == num_experts
+    assert indptr.ndim == 1
+    assert indptr.shape[0] == num_experts
+    assert indptr.dtype == "int64"
+
+    x_shape = x.shape
+    if x.ndim > 2:
+        x = x.reshape(m, k)
+        x_scale = x_scale.reshape(m, x_scale.shape[-1])
+
+    func_name = "cutlass.groupwise_scaled_group_gemm_e4m3fn_e4m3fn"
+    workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace")
+    out = op.extern(
+        func_name,
+        args=[
+            x,
+            weight,
+            x_scale,
+            weight_scale,
+            indptr,
+            workspace,
+            block_size[0],
+            block_size[1],
+        ],
+        out=nn.Tensor.placeholder((m, n), dtype=out_dtype),
+    )
+    return out.reshape(*x_shape[:-1], n) if len(x_shape) > 2 else out

python/mlc_llm/op/extern.py

@@ -0,0 +1,76 @@
+"""Potential externel modules managed by MLC compilation stack.
+
+An externl module could contain one or multiple handcrafted kernels, as long as it is provided as
+an object file (`.o`), a C++ source file (`.cc`), or a CUDA source file (`.cu`). It can be
+integrated into the system pretty smoothly.
+
+As examples, `flashinfer.py` contains such an example that instructs MLC to compile
+"$tvm_home/3rdparty/flashinfer/src/tvm_wrapper.cu" with a specific set of compilation flags and then
+link into the generated artifact of MLC LLM. TVM PR #16247
+(https://github.com/apache/tvm/pull/16247/) provides more details of using TVM's
+`nn.SourceModule` to integrate C++ and CUDA files, and `nn.ObjectModule` to integrate object files.
+
+To conveniently use those externel modules, MLC LLM compilation pipeline manages an extra global
+singleton `Store: ExternalModuleStore` to store the configured modules. It is supposed to be enabled
+before any compilation happens, and configured during a model's `forward` method is invoked.
+"""
+
+import dataclasses
+from typing import Optional
+
+from tvm.target import Target
+
+
+@dataclasses.dataclass
+class ExternModuleStore:
+    """Global store of external modules enabled during compilation."""
+
+    configured: bool = False
+    target: Optional[Target] = None
+    flashinfer: bool = False
+    faster_transformer: bool = False
+    cutlass_group_gemm: bool = False
+    cutlass_gemm: bool = False
+
+
+STORE: ExternModuleStore = ExternModuleStore()
+"""Singleton of `ExternModuleStore`."""
+
+
+def enable(target: Target, flashinfer: bool, faster_transformer: bool, cutlass: bool) -> None:
+    """Enable external modules. It should be called before any compilation happens."""
+    global STORE
+    cutlass = (
+        cutlass
+        and target.kind.name == "cuda"
+        and target.attrs.get("arch", "") in ["sm_90a", "sm_100a"]
+    )
+    faster_transformer = False
+    STORE = ExternModuleStore(
+        configured=False,
+        target=target,
+        flashinfer=flashinfer,
+        faster_transformer=faster_transformer,
+        cutlass_group_gemm=cutlass,
+        cutlass_gemm=cutlass,
+    )
+
+
+def get_store() -> ExternModuleStore:
+    """Get the global store of external modules."""
+    return STORE
+
+
+def configure() -> None:
+    """Configure external modules with extra parameters. It should be called during a model's
+    `forward` method is invoked.
+
+    Parameters
+    ----------
+    """
+    store = get_store()
+    if store.configured:
+        return
+    store.configured = True
+    if store.flashinfer or store.faster_transformer:
+        assert store.target.kind.name == "cuda"

python/mlc_llm/op/ft_gemm.py

@@ -0,0 +1,134 @@
+"""Operators enabled by external modules."""
+
+import operator
+from functools import reduce
+from typing import Optional
+
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import op
+
+
+def faster_transformer_dequantize_gemm(
+    x: nn.Tensor,
+    weight: nn.Tensor,
+    scale: nn.Tensor,
+    bias: Optional[nn.Tensor] = None,
+    activation: Optional[str] = None,
+    group_size: Optional[int] = None,
+):
+    """
+    Faster Transformer dequantize gemm inference with CutlassFpAIntB
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [*m, k].
+
+    weight : nn.Tensor
+        The quantized weight data tensor, with shape of [k, n // num_elem_per_storage].
+
+    scale : nn.Tensor
+        The quantized weight scale tensor, with shape of [k // group_size, n].
+
+    bias : Optional[nn.Tensor]
+        The optional bias for matmul, with shape broadcastable to [*m, n].
+
+    group_size : Optional[int]
+        The optional group size. If not set, then using k as group size.
+
+    Returns
+    ------
+    ret: nn.Tensor
+        The output tensor of deocde matmul, with shape of [*m, n].
+    """
+    assert x.dtype == "float16" and x.ndim >= 1
+    assert weight.ndim == 2
+    assert scale.dtype == "float16" and scale.ndim == 2
+    assert x.shape[-1] == weight.shape[0], (
+        f"Reduction dimension mismatched between x and weight, {x.shape[-1]} vs {weight.shape[0]}."
+    )
+    assert activation in [
+        None,
+        "relu",
+        "gelu",
+        "silu",
+        "identity",
+    ], "Supported activations are [None, 'identity', 'gelu', 'silu', 'relu']."
+    activation = activation if activation else "identity"
+    m = reduce(operator.mul, x.shape[:-1], 1)
+    k = x.shape[-1]
+    n = scale.shape[1]
+
+    if not group_size:
+        group_size = k
+
+    if bias:
+        assert bias.dtype == "float16" and bias.ndim >= 1
+        bias_stride = (
+            bias.shape[-1]
+            if bias and not reduce(operator.mul, bias.shape, 1) == bias.shape[-1]
+            else 0
+        )
+        return op.extern(
+            name="fastertransformer.gemm_fp16_int_bias",
+            args=[
+                x,
+                weight,
+                scale,
+                bias,
+                activation,
+                m,
+                n,
+                k,
+                group_size,
+                bias_stride,
+            ],
+            out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"),
+        )
+    return op.extern(
+        name="fastertransformer.gemm_fp16_int",
+        args=[x, weight, scale, activation, m, n, k, group_size],
+        out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"),
+    )
+
+
+def faster_transformer_moe_gemm(
+    x: nn.Tensor,
+    weight: nn.Tensor,
+    total_rows_before: nn.Tensor,
+):
+    """
+    Faster Transformer moe gemm inference with CutlassFpAIntB
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [*m, k].
+
+    weight : nn.Tensor
+        The weight data tensor, with shape of [num_experts, n, k].
+
+    total_rows_before : nn.Tensor
+        The total rows before tensor the current expert, with shape of [num_experts]. This is the
+        same as the indptr excluding the first zero element.
+
+    Returns
+    ------
+    ret: nn.Tensor
+        The output tensor of deocde matmul, with shape of [*m, n].
+    """
+    assert x.dtype == "float16" and x.ndim >= 1
+    assert weight.dtype == "float16" and weight.ndim == 3
+    assert x.shape[-1] == weight.shape[-1], (
+        f"Reduction dimension mismatched between x and weight, {x.shape[-1]} vs {weight.shape[-1]}."
+    )
+    m = reduce(operator.mul, x.shape[:-1], 1)
+    num_experts = weight.shape[0]
+    n = weight.shape[1]
+    k = x.shape[-1]
+
+    return op.extern(
+        name="fastertransformer.moe_gemm_fp16_fp16",
+        args=[x, weight, total_rows_before, m, n, k, num_experts],
+        out=nn.Tensor.placeholder((*x.shape[:-1], n), dtype="float16"),
+    )

python/mlc_llm/op/moe_matmul.py

@@ -0,0 +1,768 @@
+"""Mixture of Experts operators"""
+
+from typing import Literal, Optional, Tuple  # noqa: UP035
+
+from tvm import DataType, DataTypeCode, s_tir, tirx
+from tvm.relax.frontend.nn import Tensor, op
+from tvm.script import tirx as T
+
+# mypy: disable-error-code="attr-defined,valid-type,name-defined"
+
+
+def gemv(x: Tensor, w: Tensor, indptr: Tensor) -> Tensor:
+    """GEMV for project-in (e1-e3) or project-out (e2) in MLP.
+
+    Parameters
+    ----------
+    x : Tensor
+        For project-in, the input tensor of shape (1, in_features); and for project-out, the input
+        shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
+        experts per token.
+
+    w : Tensor
+        The weight tensor of shape (local_experts, out_features, in_features), where `local_experts`
+        is the total number of experts.
+
+    indptr : Tensor
+        The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
+        number of activated experts per token.
+
+    Returns
+    -------
+    out : Tensor
+        The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the
+        number of activated experts per token.
+    """
+    (local_experts, out_features, in_features), dtype = w.shape, w.dtype
+    _, experts_per_tok = indptr.shape
+    x_leading_dim, _ = x.shape
+
+    def access_x(x, e, j):
+        return x[0, j] if x_leading_dim == 1 else x[e, j]
+
+    # NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily.
+    assert w.shape == [local_experts, out_features, in_features] and w.dtype == dtype
+    assert x.shape == [x_leading_dim, in_features] and x.dtype == dtype
+    assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32"
+    assert x_leading_dim in [1, experts_per_tok]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        x: T.Buffer((x_leading_dim, in_features), dtype),
+        w: T.Buffer((local_experts, out_features, in_features), dtype),
+        indptr: T.Buffer((1, experts_per_tok), "int32"),
+        o: T.Buffer((experts_per_tok, out_features), dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        for e in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
+            with T.sblock("gemv_o"):
+                e = T.axis.spatial(experts_per_tok, e)
+                T.reads(x[:, :], w[indptr[0, e], :, :], indptr[0, e])
+                T.writes(o[e, :])
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("gemv"):
+                        i, j = T.axis.remap("SR", [i1, i2])
+                        with T.init():
+                            o[e, i] = T.cast(T.float16(0), dtype)
+                        o[e, i] += access_x(x, e, j) * w[indptr[0, e], i, j]
+
+    return op.tensor_ir_op(
+        _func,
+        "moe_gemv",
+        args=[x, w, indptr],
+        out=Tensor.placeholder([experts_per_tok, out_features], dtype),
+    )
+
+
+def dequantize_gemv(
+    x: Tensor,
+    w: Tensor,
+    scale: Tensor,
+    indptr: Tensor,
+    quantize_dtype: str,
+    group_size: int,
+) -> Tensor:
+    """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized.
+    It needs to be dequantized before the GEMV computation.
+
+    Parameters
+    ----------
+    x : Tensor
+        For project-in, the input tensor of shape (1, in_features); and for project-out, the input
+        shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
+        experts per token.
+
+    w : Tensor
+        The quantized weight tensor of shape (local_experts, out_features, in_features // n),
+        where n is the number of elements per storage dtype, e.g. if the storage dtype is uint32,
+        and the quantize dtype is int4, then n is 8.
+        `local_experts` is the total number of experts including activated and non-active ones.
+
+    scale : Tensor
+        The scale tensor of shape (local_experts, out_features, in_features // group_size), where
+        `local_experts` is the total number of experts including activated and non-active ones.
+
+    indptr : Tensor
+        The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
+        number of activated experts per token.
+
+    quantize_dtype : str
+        The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc.
+
+    group_size : int
+        The number of elements in each quantization group, e.g. 32 or 128.
+
+    Returns
+    -------
+    out : Tensor
+        The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the
+        number of activated experts per token.
+    """
+    (x_leading_dim, in_features), model_dtype = x.shape, x.dtype
+    (local_experts, out_features, _), storage_dtype = w.shape, w.dtype
+    _, experts_per_tok = indptr.shape
+    quantize_dtype_bits = DataType(quantize_dtype).bits
+    num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
+    num_group = (in_features + group_size - 1) // group_size
+    num_storage = group_size // num_elem_per_storage * num_group
+
+    def _dequantize(w, s, e, i, j):
+        tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype)
+        tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype)
+        w = w[e, i, j // num_elem_per_storage]
+        s = s[e, i, j // group_size]
+        shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
+        w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype)
+        return (w - tir_max_int) * s
+
+    def access_x(x, e, j):
+        return x[0, j] if x_leading_dim == 1 else x[e, j]
+
+    assert x.shape == [x_leading_dim, in_features] and x.dtype == model_dtype
+    assert w.shape == [local_experts, out_features, num_storage] and w.dtype == storage_dtype
+    assert scale.shape == [local_experts, out_features, num_group] and scale.dtype == model_dtype
+    assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32"
+    assert x_leading_dim in [1, experts_per_tok]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        x: T.Buffer((x_leading_dim, in_features), model_dtype),
+        w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
+        scale: T.Buffer((local_experts, out_features, num_group), model_dtype),
+        indptr: T.Buffer((1, experts_per_tok), "int32"),
+        o: T.Buffer((experts_per_tok, out_features), model_dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
+            with T.sblock("gemv_o"):
+                e = T.axis.spatial(experts_per_tok, expert_id)
+                y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("dequantize"):
+                        i, j = T.axis.remap("SS", [i1, i2])
+                        y[i, j] = _dequantize(w, scale, indptr[0, e], i, j)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("gemv"):
+                        i, j = T.axis.remap("SR", [i1, i2])
+                        with T.init():
+                            o[e, i] = T.cast(T.float16(0), model_dtype)
+                        o[e, i] += access_x(x, e, j) * y[i, j]
+
+    return op.tensor_ir_op(
+        _func,
+        "moe_dequantize_gemv",
+        args=[x, w, scale, indptr],
+        out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
+    )
+
+
+def dequantize_float8_gemv(
+    x: Tensor,
+    w: Tensor,
+    scale: Optional[Tensor],
+    indptr: Tensor,
+    quantize_dtype: Literal["float8_e5m2", "float8_e4m3fn"],
+) -> Tensor:
+    """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in
+    fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation.
+
+    Parameters
+    ----------
+    x : Tensor
+        For project-in, the input tensor of shape (1, in_features); and for project-out, the input
+        shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
+        experts per token.
+
+    w : Tensor
+        The quantized weight tensor of shape (local_experts, out_features, in_features)
+
+    scale : Optional[Tensor]
+        The optional scale tensor of shape (1,)
+
+    indptr : Tensor
+        The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
+        number of activated experts per token.
+
+    quantize_dtype : Literal["float8_e5m2", "float8_e4m3fn"]
+        The quantize dtype of the weight tensor, which is either float8_e5m2 or float8_e4m3fn.
+    """
+    (x_leading_dim, in_features), model_dtype = x.shape, x.dtype
+    (local_experts, out_features, _), storage_dtype = w.shape, w.dtype
+    _, experts_per_tok = indptr.shape
+    quantize_dtype_bits = DataType(quantize_dtype).bits
+    num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
+    num_storage = tirx.ceildiv(in_features, num_elem_per_storage)
+
+    def _dequantize(w, s, e, i, j):
+        if num_elem_per_storage == 1:
+            w = tirx.reinterpret(quantize_dtype, w[e, i, j])
+        else:
+            assert DataType(storage_dtype).type_code == DataTypeCode.UINT
+            tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype)
+            w = w[e, i, j // num_elem_per_storage]
+            shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
+            w = tirx.reinterpret(
+                quantize_dtype,
+                tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype("uint8"),
+            )
+        w = w.astype(model_dtype)
+        if s is not None:
+            w = w * s[0]
+        return w
+
+    def access_x(x, e, j):
+        return x[0, j] if x_leading_dim == 1 else x[e, j]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func_with_scale(
+        x: T.Buffer((x_leading_dim, in_features), model_dtype),
+        w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
+        scale: T.Buffer((1,), "float32"),
+        indptr: T.Buffer((1, experts_per_tok), "int32"),
+        o: T.Buffer((experts_per_tok, out_features), model_dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
+            with T.sblock("gemv_o"):
+                e = T.axis.spatial(experts_per_tok, expert_id)
+                y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("dequantize"):
+                        i, j = T.axis.remap("SS", [i1, i2])
+                        y[i, j] = _dequantize(w, scale, indptr[0, e], i, j)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("gemv"):
+                        i, j = T.axis.remap("SR", [i1, i2])
+                        with T.init():
+                            o[e, i] = T.cast(T.float16(0), model_dtype)
+                        o[e, i] += access_x(x, e, j) * y[i, j]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func_without_scale(
+        x: T.Buffer((x_leading_dim, in_features), model_dtype),
+        w: T.Buffer((local_experts, out_features, num_storage), storage_dtype),
+        indptr: T.Buffer((1, experts_per_tok), "int32"),
+        o: T.Buffer((experts_per_tok, out_features), model_dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
+            with T.sblock("gemv_o"):
+                e = T.axis.spatial(experts_per_tok, expert_id)
+                y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("dequantize"):
+                        i, j = T.axis.remap("SS", [i1, i2])
+                        y[i, j] = _dequantize(w, None, indptr[0, e], i, j)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("gemv"):
+                        i, j = T.axis.remap("SR", [i1, i2])
+                        with T.init():
+                            o[e, i] = T.cast(T.float16(0), model_dtype)
+                        o[e, i] += access_x(x, e, j) * y[i, j]
+
+    if scale is not None:
+        return op.tensor_ir_op(
+            _func_with_scale,
+            "moe_dequantize_gemv",
+            args=[x, w, scale, indptr],
+            out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
+        )
+    return op.tensor_ir_op(
+        _func_without_scale,
+        "moe_dequantize_gemv",
+        args=[x, w, indptr],
+        out=Tensor.placeholder([experts_per_tok, out_features], model_dtype),
+    )
+
+
+def dequantize_block_scale_float8_gemv(
+    x: Tensor,
+    w: Tensor,
+    w_scale: Tensor,
+    expert_indices: Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+) -> Tensor:
+    """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in
+    fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation.
+
+    Parameters
+    ----------
+    x : Tensor
+        For project-in, the input tensor of shape (1, in_features); and for project-out, the input
+        shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated
+        experts per token.
+
+    w : Tensor
+        The quantized weight tensor of shape (local_experts, out_features, in_features)
+
+    w_scale : Tensor
+        The scale tensor of shape
+        (local_experts, out_features // block_size[0], in_features // block_size[1])
+
+    indptr : Tensor
+        The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the
+        number of activated experts per token.
+
+    block_size : Tuple[int, int]
+        The block size of the weight tensor.
+
+    out_dtype : str
+        The output dtype of the GEMV computation.
+    """
+    x_leading_dim, in_features = x.shape
+    local_experts, out_features, k = w.shape
+    _, experts_per_tok = expert_indices.shape
+    model_dtype = x.dtype
+    quantize_dtype = w.dtype
+
+    assert out_features % block_size[0] == 0
+    assert k % block_size[1] == 0
+
+    def _dequantize(w, s, e, i, j):
+        return w[e, i, j].astype(model_dtype) * s[e, i // block_size[0], j // block_size[1]].astype(
+            model_dtype
+        )
+
+    def load_x(x, e, j):
+        return x[0, j] if x_leading_dim == 1 else x[e, j]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        x: T.Buffer((x_leading_dim, in_features), model_dtype),
+        w: T.Buffer((local_experts, out_features, k), quantize_dtype),
+        w_scale: T.Buffer(
+            (local_experts, out_features // block_size[0], k // block_size[1]),
+            "float32",
+        ),
+        expert_indices: T.Buffer((1, experts_per_tok), "int32"),
+        o: T.Buffer((experts_per_tok, out_features), out_dtype),
+    ):
+        T.func_attr({"op_pattern": 4, "tirx.noalias": True})  # kOutEWiseFusable
+        for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"):
+            with T.sblock("gemv_o"):
+                e = T.axis.spatial(experts_per_tok, expert_id)
+                y = T.sblock_alloc_buffer((out_features, in_features), model_dtype)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("dequantize"):
+                        i, j = T.axis.remap("SS", [i1, i2])
+                        y[i, j] = _dequantize(w, w_scale, expert_indices[0, e], i, j)
+                for i1, i2 in T.grid(out_features, in_features):
+                    with T.sblock("gemv"):
+                        i, j = T.axis.remap("SR", [i1, i2])
+                        with T.init():
+                            o[e, i] = T.cast(T.float16(0), out_dtype)
+                        o[e, i] += (load_x(x, e, j) * y[i, j]).astype(out_dtype)
+
+    return op.tensor_ir_op(
+        _func,
+        "moe_dequantize_gemv",
+        args=[x, w, w_scale, expert_indices],
+        out=Tensor.placeholder([experts_per_tok, out_features], out_dtype),
+    )
+
+
+def group_gemm(x: Tensor, w: Tensor, indptr: Tensor):
+    """Group GEMM in MoE models.
+
+    Parameters
+    ----------
+    x : Tensor
+        Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape.
+
+    w : Tensor
+        Weight tensor of shape (num_local_experts, out_features, in_features).
+        `w[i, :, :]` is the weight matrix for the `i`-th local expert.
+
+    indptr : Tensor
+        Index pointer tensor of shape (num_local_experts + 1, ).
+        `x[indptr[a] : indptr[a + 1]]` is the input for the `i`-th local expert.
+
+    Returns
+    -------
+    out : Tensor
+        Output tensor of shape (batch_size, out_features).
+    """
+    # NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily.
+    (num_local_experts, out_features, in_features), dtype = w.shape, w.dtype
+
+    assert x.shape[1:] == [in_features] and x.dtype == dtype
+    assert indptr.shape == [num_local_experts + 1] and indptr.dtype == "int32"
+
+    Ne, N, K = num_local_experts, out_features, in_features
+    BLK_M, BLK_N, BLK_K = 8, 128, 32
+    TX, TY, CTA_COUNT = 8, 32, 1024
+    VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1
+    UNROLL = 64
+    STORAGE_ALIGN = False
+    assert BLK_K % 8 == 0
+    tiles_per_row = (N + BLK_N - 1) // BLK_N
+    zero = tirx.const(0, dtype)
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_x: T.handle,
+        var_w: T.handle,
+        var_indptr: T.handle,
+        var_o: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        B = T.int32()
+        X = T.match_buffer(var_x, (B, K), dtype)
+        W = T.match_buffer(var_w, (Ne, N, K), dtype)
+        indptr = T.match_buffer(var_indptr, (Ne + 1,), "int32")
+        out = T.match_buffer(var_o, (B, N), dtype)
+
+        for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"):
+            with T.sblock("CTA"):
+                bx = T.axis.spatial(CTA_COUNT, _bx)
+                T.reads(indptr[:], X[:, :], W[:, :, :])
+                T.writes(out[:, :])
+                sum = T.sblock_alloc_buffer((2,), "int32", scope="local")
+                row = T.sblock_alloc_buffer((2,), "int32", scope="local")
+                cur_e = T.sblock_alloc_buffer((1,), "int32", scope="local")
+                tile_id = T.sblock_alloc_buffer((1,), "int32", scope="local")
+                sum[0] = 0
+                sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row
+                row[0] = 0
+                row[1] = indptr[1] - indptr[0]
+                cur_e[0] = 0
+                tile_id[0] = bx
+                while T.tvm_thread_invariant(cur_e[0] < Ne):
+                    # move to the current group
+                    while sum[1] <= tile_id[0] and cur_e[0] < Ne:
+                        cur_e[0] += 1
+                        if cur_e[0] < Ne:
+                            e: T.int32 = cur_e[0]
+                            delta: T.int32 = indptr[e + 1] - indptr[e]
+                            sum[0] = sum[1]
+                            sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row
+                            row[0] = row[1]
+                            row[1] += delta
+                    # sync threads to make sure all threads have the same tile position
+                    T.tvm_storage_sync("shared")
+                    if T.tvm_thread_invariant(cur_e[0] < Ne):
+                        # fetch current tile position
+                        e: T.int32 = cur_e[0]
+                        num_tiles: T.int32 = tile_id[0] - sum[0]
+                        m_offset: T.int32 = BLK_M * T.floordiv(num_tiles, tiles_per_row) + row[0]
+                        n_offset: T.int32 = BLK_N * T.floormod(num_tiles, tiles_per_row)
+                        with T.sblock("gemm"):
+                            T.reads(
+                                row[1],
+                                X[m_offset : m_offset + BLK_M, :],
+                                W[e, n_offset : n_offset + BLK_N, :],
+                            )
+                            T.writes(
+                                out[
+                                    m_offset : m_offset + BLK_M,
+                                    n_offset : n_offset + BLK_N,
+                                ]
+                            )
+                            X_tile = T.sblock_alloc_buffer((BLK_M, K), dtype, scope="shared")
+                            W_tile = T.sblock_alloc_buffer((BLK_N, K), dtype, scope="shared")
+                            O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), dtype, scope="local")
+                            for a0, a1 in T.grid(BLK_M, K):
+                                with T.sblock("X_shared"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    X_tile[i, j] = T.if_then_else(
+                                        m_offset + i < row[1],
+                                        X[m_offset + i, j],
+                                        zero,
+                                    )
+                            for a0, a1 in T.grid(BLK_N, K):
+                                with T.sblock("W_shared"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    W_tile[i, j] = T.if_then_else(
+                                        n_offset + i < N,
+                                        W[e, n_offset + i, j],
+                                        zero,
+                                    )
+                            for a0, a1, a2 in T.grid(BLK_M, BLK_N, K):
+                                with T.sblock("compute"):
+                                    i, j, k = T.axis.remap("SSR", [a0, a1, a2])
+                                    with T.init():
+                                        O_tile[i, j] = zero
+                                    O_tile[i, j] += X_tile[i, k] * W_tile[j, k]
+                            for a0, a1 in T.grid(BLK_M, BLK_N):
+                                with T.sblock("store"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    if m_offset + i < row[1] and n_offset + j < N:
+                                        out[m_offset + i, n_offset + j] = O_tile[i, j]
+                    # move to next tile
+                    tile_id[0] += CTA_COUNT
+
+    def _schedule():
+        sch = s_tir.Schedule(_func)
+
+        def _cooperative_fetch(block, vec_len):
+            num_loops = len(sch.get_loops(block))
+            sch.compute_at(block, ko, preserve_unit_loops=True)
+            loops = sch.get_loops(block)[-num_loops:]
+            ty, tx, _, vec = sch.split(
+                sch.fuse(*loops),
+                factors=[TY, TX, None, vec_len],
+            )
+            sch.vectorize(vec)
+            sch.bind(ty, "threadIdx.y")
+            sch.bind(tx, "threadIdx.x")
+            if STORAGE_ALIGN:
+                sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len)
+            return block
+
+        main_block = sch.get_sblock("compute")
+        x, y, k = sch.get_loops(main_block)
+        ty, yi = sch.split(y, [TY, None])
+        tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT])
+        ko, ki = sch.split(k, factors=[None, BLK_K])
+        sch.reorder(ty, tx, ko, ki, yi, xi, vec_c)
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.vectorize(vec_c)
+        if UNROLL > 0:
+            sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL)
+            sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+        l2g = sch.get_sblock("store")
+        sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
+        _, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O])
+        sch.vectorize(v)
+        _cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X)
+        _cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W)
+        sch.decompose_reduction(main_block, ko)
+        return sch.mod["main"]
+
+    return op.tensor_ir_op(
+        _schedule(),
+        "group_gemm",
+        args=[x, w, indptr],
+        out=Tensor.placeholder([x.shape[0], out_features], dtype),
+    )
+
+
+def dequantize_group_gemm(
+    x: Tensor,
+    w: Tensor,
+    scale: Tensor,
+    indptr: Tensor,
+    quantize_dtype: str,
+    indptr_dtype: str,
+    group_size: int,
+):
+    """Group GEMM in MoE models but the weight is quantized.
+
+    Parameters
+    ----------
+    x : Tensor
+        Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape.
+
+    w : Tensor
+        Weight tensor of shape (num_local_experts, out_features, in_features // n), where n is the
+        number of elements per storage dtype, e.g. if the storage dtype is uint32, and the quantize
+        dtype is int4, then n is 8.
+
+    scale : Tensor
+        The scale tensor of shape (num_local_experts, out_features, in_features // group_size).
+
+    indptr : Tensor
+        Index pointer tensor of shape (num_local_experts + 1, ). `x[indptr[a] : indptr[a + 1]]` is
+        the input for the `i`-th local expert.
+
+    group_size : int
+        The number of elements in each quantization group, e.g. 32 or 128.
+
+    quantize_dtype : str
+        The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc.
+
+    indptr_dtype : str
+        The dtype of the index pointer tensor, which can be int32 or int64.
+
+    Returns
+    -------
+    out : Tensor
+        Output tensor of shape (batch_size, out_features).
+    """
+    (_, in_features), model_dtype = x.shape, x.dtype
+    (num_local_experts, out_features, _), storage_dtype = w.shape, w.dtype
+    quantize_dtype_bits = DataType(quantize_dtype).bits
+    num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits
+    num_group = (in_features + group_size - 1) // group_size
+    num_storage = group_size // num_elem_per_storage * num_group
+
+    def _dequantize(w, s, e, i, j):
+        tir_bin_mask = tirx.const((1 << quantize_dtype_bits) - 1, storage_dtype)
+        tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype)
+        w = w[e, i, j // num_elem_per_storage]
+        s = s[e, i, j // group_size]
+        shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype)
+        w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype)
+        return (w - tir_max_int) * s
+
+    Ne, N, K = num_local_experts, out_features, in_features
+    BLK_M, BLK_N, BLK_K = 8, 128, 32
+    TX, TY, CTA_COUNT = 8, 32, 1024
+    VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1
+    UNROLL = 64
+    STORAGE_ALIGN = False
+    assert BLK_K % 8 == 0
+    tiles_per_row = (N + BLK_N - 1) // BLK_N
+    zero = tirx.const(0, model_dtype)
+    if indptr_dtype == "int64":
+        indptr = op.pad(indptr, [1, 0], "constant", 0)
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_x: T.handle,
+        w: T.Buffer((Ne, N, num_storage), storage_dtype),
+        scale: T.Buffer((Ne, N, num_group), model_dtype),
+        indptr: T.Buffer((Ne + 1,), indptr_dtype),
+        var_o: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        B = T.int32()
+        X = T.match_buffer(var_x, (B, K), model_dtype)
+        out = T.match_buffer(var_o, (B, N), model_dtype)
+        for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"):
+            with T.sblock("CTA"):
+                bx = T.axis.spatial(CTA_COUNT, _bx)
+                T.reads(X[:, :], w[:, :, :], scale[:, :, :], indptr[:])
+                T.writes(out[:, :])
+                sum = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local")
+                row = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local")
+                cur_e = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local")
+                tile_id = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local")
+                sum[0] = 0
+                sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row
+                row[0] = 0
+                row[1] = indptr[1] - indptr[0]
+                cur_e[0] = 0
+                tile_id[0] = bx
+                while T.tvm_thread_invariant(cur_e[0] < Ne):
+                    # move to the current group
+                    while sum[1] <= tile_id[0] and cur_e[0] < Ne:
+                        cur_e[0] += 1
+                        if cur_e[0] < Ne:
+                            e = cur_e[0]
+                            delta = indptr[e + 1] - indptr[e]
+                            sum[0] = sum[1]
+                            sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row
+                            row[0] = row[1]
+                            row[1] += delta
+                    # sync threads to make sure all threads have the same tile position
+                    T.tvm_storage_sync("shared")
+                    if T.tvm_thread_invariant(cur_e[0] < Ne):
+                        # fetch current tile position
+                        e = cur_e[0]
+                        num_tiles = tile_id[0] - sum[0]
+                        m_offset = T.floordiv(num_tiles, tiles_per_row) * BLK_M + row[0]
+                        n_offset = T.floormod(num_tiles, tiles_per_row) * BLK_N
+                        with T.sblock("gemm"):
+                            T.reads(
+                                row[1],
+                                X[m_offset : m_offset + BLK_M, :],
+                                w[e, n_offset : n_offset + BLK_N, :],
+                                scale[e, n_offset : n_offset + BLK_N, :],
+                            )
+                            T.writes(
+                                out[
+                                    m_offset : m_offset + BLK_M,
+                                    n_offset : n_offset + BLK_N,
+                                ]
+                            )
+                            X_tile = T.sblock_alloc_buffer((BLK_M, K), model_dtype, scope="shared")
+                            W_tile = T.sblock_alloc_buffer((BLK_N, K), model_dtype, scope="shared")
+                            O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), "float32", scope="local")
+                            for a0, a1 in T.grid(BLK_M, K):
+                                with T.sblock("X_shared"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    X_tile[i, j] = T.if_then_else(
+                                        m_offset + i < row[1],
+                                        X[m_offset + i, j],
+                                        zero,
+                                    )
+                            for a0, a1 in T.grid(BLK_N, K):
+                                with T.sblock("W_shared"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    W_tile[i, j] = T.if_then_else(
+                                        n_offset + i < N,
+                                        _dequantize(w, scale, e, n_offset + i, j),
+                                        zero,
+                                    )
+                            for a0, a1, a2 in T.grid(BLK_M, BLK_N, K):
+                                with T.sblock("compute"):
+                                    i, j, k = T.axis.remap("SSR", [a0, a1, a2])
+                                    with T.init():
+                                        O_tile[i, j] = zero
+                                    O_tile[i, j] += X_tile[i, k] * W_tile[j, k]
+                            for a0, a1 in T.grid(BLK_M, BLK_N):
+                                with T.sblock("store"):
+                                    i, j = T.axis.remap("SS", [a0, a1])
+                                    if m_offset + i < row[1] and n_offset + j < N:
+                                        out[m_offset + i, n_offset + j] = O_tile[i, j]
+                    # move to next tile
+                    tile_id[0] += CTA_COUNT
+
+    def _schedule():
+        sch = s_tir.Schedule(_func)
+
+        def _cooperative_fetch(block, vec_len):
+            num_loops = len(sch.get_loops(block))
+            sch.compute_at(block, ko, preserve_unit_loops=True)
+            loops = sch.get_loops(block)[-num_loops:]
+            ty, tx, _, vec = sch.split(
+                sch.fuse(*loops),
+                factors=[TY, TX, None, vec_len],
+            )
+            sch.vectorize(vec)
+            sch.bind(ty, "threadIdx.y")
+            sch.bind(tx, "threadIdx.x")
+            if STORAGE_ALIGN:
+                sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len)
+            return block
+
+        main_block = sch.get_sblock("compute")
+        x, y, k = sch.get_loops(main_block)
+        ty, yi = sch.split(y, [TY, None])
+        tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT])
+        ko, ki = sch.split(k, factors=[None, BLK_K])
+        sch.reorder(ty, tx, ko, ki, yi, xi, vec_c)
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.vectorize(vec_c)
+        if UNROLL > 0:
+            sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL)
+            sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
+        l2g = sch.get_sblock("store")
+        sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
+        _, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O])
+        sch.vectorize(v)
+        _cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X)
+        _cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W)
+        sch.decompose_reduction(main_block, ko)
+        return sch.mod["main"]
+
+    return op.tensor_ir_op(
+        _schedule(),
+        "dequantize_group_gemm",
+        args=[x, w, scale, indptr],
+        out=Tensor.placeholder([x.shape[0], out_features], model_dtype),
+    )

python/mlc_llm/op/moe_misc.py

@@ -0,0 +1,645 @@
+"""Mixture of Experts operators"""
+
+from functools import reduce
+from typing import Literal, Optional, Tuple, Union  # noqa: UP035
+
+import numpy as np
+from tvm import te, tirx
+from tvm.relax.frontend.nn import IntExpr, Tensor, op
+from tvm.script import tirx as T
+
+# mypy: disable-error-code="attr-defined,name-defined"
+
+
+def moe_sum(x: Tensor, dim: int) -> Tensor:
+    """Compute the sum of the input tensor along the given axis. It is specialized for the MoE
+    case where `x.ndim == 3` and `x.shape[1] == num_experts_per_tok (which is 2)`.
+    """
+
+    if x.shape[1] == 1:
+        return x.reshape(x.shape[0], x.shape[2])
+
+    if x.ndim == 3 and x.shape[1] == 2:
+        return op.tensor_expr_op(
+            lambda x: te.compute(
+                (x.shape[0], x.shape[2]),
+                lambda i, j: x[i, 0, j] + x[i, 1, j],
+                name="sum_2",
+            ),
+            "sum",
+            args=[x],
+        )
+    return op.sum(x, axis=dim)
+
+
+def _gating_topk_init_local_top_k(k_val, dtype, local_top_k, local_top_k_index):
+    for t in range(k_val):
+        T.buffer_store(local_top_k, T.min_value(dtype), indices=[t])
+    for t in range(k_val):
+        T.buffer_store(local_top_k_index, t, indices=[-1])
+
+
+def _gating_topk_process_value(k_val, x, local_top_k, local_top_k_index, vi, vk):
+    if_frames = [T.If(x[vi, vk] > local_top_k[i]) for i in range(k_val)]
+    then_frames = [T.Then() for _ in range(k_val)]
+    else_frames = [T.Else() for _ in range(k_val - 1)]
+    for i in range(k_val):
+        if_frames[i].__enter__()
+        with then_frames[i]:
+            for j in range(k_val - 1, i, -1):
+                T.buffer_store(local_top_k, local_top_k[j - 1], indices=[j])
+                T.buffer_store(local_top_k_index, local_top_k_index[j - 1], indices=[j])
+            T.buffer_store(local_top_k, x[vi, vk], indices=[i])
+            T.buffer_store(local_top_k_index, vk, indices=[i])
+        if i != k_val - 1:
+            else_frames[i].__enter__()
+
+    for i in range(k_val - 1, -1, -1):
+        if i != k_val - 1:
+            else_frames[i].__exit__(None, None, None)
+        if_frames[i].__exit__(None, None, None)
+
+
+def gating_topk(scores: Tensor, k: int) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+    """Compute the top-k experts and their scores.
+
+    Parameters
+    ----------
+    scores : Tensor
+        The input tensor with shape [batch_size, num_local_experts].
+
+    k : int
+        The number of top elements to be selected, which is `num_experts_per_tok` in MoE.
+
+    Returns
+    -------
+    expert_weights: Tensor
+        The top-k expert scores with shape [batch_size, k].
+
+    expert_indices: Tensor
+        The top-k expert indices with shape [batch_size, k].
+    """
+    (batch_size, num_local_experts), dtype = scores.shape, scores.dtype
+    index_dtype = "int32"
+
+    TX = 1024
+
+    def _get_topk_func(k_val: int):
+        @T.prim_func(private=True, s_tir=True)
+        def topk_func(
+            var_x: T.handle,
+            var_out: T.handle,
+            var_out_index: T.handle,
+        ) -> None:
+            T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": True})
+            batch_size = T.int64()
+            x = T.match_buffer(var_x, (batch_size, num_local_experts), dtype)
+            out = T.match_buffer(var_out, (batch_size, k_val), dtype)
+            out_index = T.match_buffer(var_out_index, (batch_size, k_val), index_dtype)
+            local_top_k = T.sblock_alloc_buffer((k_val,), dtype=dtype, scope="local")
+            local_top_k_index = T.sblock_alloc_buffer((k_val,), dtype=index_dtype, scope="local")
+            for io in T.thread_binding(0, T.ceildiv(batch_size, TX), "blockIdx.x"):
+                for ii in T.thread_binding(0, TX, "threadIdx.x"):
+                    with T.sblock("top_k"):
+                        vi = T.axis.spatial(batch_size, io * TX + ii)
+                        T.where(io * TX + ii < batch_size)
+                        with T.sblock("init"):
+                            _gating_topk_init_local_top_k(
+                                k_val, dtype, local_top_k, local_top_k_index
+                            )
+                        for k in range(num_local_experts):
+                            with T.sblock("update"):
+                                vk = T.axis.remap("S", [k])
+                                _gating_topk_process_value(
+                                    k_val, x, local_top_k, local_top_k_index, vi, vk
+                                )
+                        for j in T.unroll(k_val):
+                            with T.sblock("output"):
+                                vj = T.axis.remap("S", [j])
+                                out[vi, vj] = local_top_k[vj]
+                                out_index[vi, vj] = local_top_k_index[vj]
+
+        return topk_func
+
+    return op.tensor_ir_op(
+        _get_topk_func(k),
+        f"top{k}",
+        args=[scores],
+        out=(
+            Tensor.placeholder([batch_size, k], dtype),
+            Tensor.placeholder([batch_size, k], index_dtype),
+        ),
+    )
+
+
+def gating_softmax_topk(x: Tensor, k: int, norm_topk_prob=True) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+    """Compute the softmax score, choose the top-k experts, and returns selected scores.
+
+    Parameters
+    ----------
+    x : Tensor
+        The input tensor with shape [batch_size, num_local_experts].
+
+    k : int
+        The number of top elements to be selected, which is `num_experts_per_tok` in MoE.
+
+    norm_topk_prob : bool
+        Whether to normalize the top-k expert scores.
+
+    Returns
+    -------
+    expert_weights: Tensor
+        The top-k expert scores with shape [batch_size, k].
+
+    expert_indices: Tensor
+        The top-k expert indices with shape [batch_size, k].
+    """
+    (batch_size, num_local_experts), dtype = x.shape, x.dtype
+    index_dtype = "int32"
+
+    TX = 1024
+
+    def _get_topk_softmax_norm_func(k_val: int):
+        def _nested_max(local_top_k_f32):
+            expr = local_top_k_f32[0]
+            for i in range(1, k_val):
+                expr = T.max(expr, local_top_k_f32[i])
+            return expr
+
+        def _nested_sum(local_top_k_f32, local_top_k_max):
+            expr = T.exp(local_top_k_f32[0] - local_top_k_max[0])
+            for i in range(1, k_val):
+                expr = expr + T.exp(local_top_k_f32[i] - local_top_k_max[0])
+            return expr
+
+        @T.prim_func(private=True, s_tir=True)
+        def topk_softmax_norm_func(
+            var_x: T.handle,
+            var_out: T.handle,
+            var_out_index: T.handle,
+        ) -> None:
+            T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": True})
+            batch_size = T.int64()
+            x = T.match_buffer(var_x, (batch_size, num_local_experts), dtype)
+            out = T.match_buffer(var_out, (batch_size, k_val), dtype)
+            out_index = T.match_buffer(var_out_index, (batch_size, k_val), index_dtype)
+            local_top_k = T.sblock_alloc_buffer((k_val,), dtype=dtype, scope="local")
+            local_top_k_index = T.sblock_alloc_buffer((k_val,), dtype=index_dtype, scope="local")
+            local_top_k_f32 = T.sblock_alloc_buffer((k_val,), dtype="float32", scope="local")
+            local_top_k_max = T.sblock_alloc_buffer((1,), dtype="float32", scope="local")
+            for io in T.thread_binding(0, T.ceildiv(batch_size, TX), "blockIdx.x"):
+                for ii in T.thread_binding(0, TX, "threadIdx.x"):
+                    with T.sblock("top_k"):
+                        vi = T.axis.spatial(batch_size, io * TX + ii)
+                        T.where(io * TX + ii < batch_size)
+                        with T.sblock("init"):
+                            _gating_topk_init_local_top_k(
+                                k_val, dtype, local_top_k, local_top_k_index
+                            )
+                        for k in range(num_local_experts):
+                            with T.sblock("update"):
+                                vk = T.axis.remap("S", [k])
+                                _gating_topk_process_value(
+                                    k_val, x, local_top_k, local_top_k_index, vi, vk
+                                )
+                        for j in T.unroll(k_val):
+                            with T.sblock("cast"):
+                                vj = T.axis.remap("S", [j])
+                                local_top_k_f32[vj] = T.cast(local_top_k[vj], "float32")
+                        with T.sblock("max"):
+                            local_top_k_max[0] = _nested_max(local_top_k_f32)
+                        for j in T.unroll(k_val):
+                            with T.sblock("output"):
+                                vj = T.axis.remap("S", [j])
+                                out[vi, vj] = T.cast(
+                                    T.exp(local_top_k_f32[vj] - local_top_k_max[0])
+                                    / _nested_sum(local_top_k_f32, local_top_k_max),
+                                    dtype,
+                                )
+                                out_index[vi, vj] = local_top_k_index[vj]
+
+        return topk_softmax_norm_func
+
+    if norm_topk_prob:
+        return op.tensor_ir_op(
+            _get_topk_softmax_norm_func(k),
+            f"top{k}_softmax",
+            args=[x],
+            out=(
+                Tensor.placeholder([batch_size, k], dtype),
+                Tensor.placeholder([batch_size, k], index_dtype),
+            ),
+        )
+
+    expert_score = op.softmax(x.astype("float32"), axis=-1).astype(dtype)
+    return gating_topk(expert_score, k)
+
+
+def group_limited_greedy_topk(
+    scores: Tensor,  # (num_tokens, num_routed_experts)
+    top_k: int,
+    num_routed_experts: int,
+    n_group: int,
+    topk_group: int,
+    topk_method: Literal["group_limited_greedy", "noaux_tc"],
+    num_tokens: IntExpr,
+    e_score_correction_bias: Optional[Tensor],
+) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+    """Group-limited greedy top-k expert selection.
+
+    Parameters
+    ----------
+    scores : Tensor
+        The input tensor with shape [num_tokens, num_routed_experts].
+
+    top_k : int
+        The number of top elements to be selected, which is `num_experts_per_tok` in MoE.
+
+    num_routed_experts : int
+        The number of routed experts.
+
+    n_group : int
+        The number of groups.
+
+    topk_group : int
+        The number of top-k groups to be selected.
+
+    topk_method : Literal["group_limited_greedy", "noaux_tc"]
+        The method to select the top-k groups.
+
+    num_tokens : IntExpr
+        The number of tokens.
+
+    e_score_correction_bias : Optional[Tensor]
+        The bias of the expert scores. Only available for "noaux_tc".
+
+    Returns
+    -------
+    expert_weights : Tensor
+        The top-k expert scores with shape [num_tokens, top_k].
+
+    expert_indices : Tensor
+        The top-k expert indices with shape [num_tokens, top_k].
+    """
+    assert scores.dtype == "float32"
+    scores_for_choice = scores
+    if topk_method == "noaux_tc":
+        assert e_score_correction_bias is not None
+        assert e_score_correction_bias.dtype == "float32"
+        scores_for_choice = scores + e_score_correction_bias
+    group_size = num_routed_experts // n_group
+    if topk_method == "noaux_tc":
+        group_scores = op.sum(
+            gating_topk(
+                scores_for_choice.reshape(num_tokens * n_group, group_size),
+                2,
+            )[0],
+            axis=-1,
+        ).reshape(num_tokens, n_group)
+    else:
+        group_scores = op.max(
+            scores_for_choice.reshape(num_tokens * n_group, group_size), axis=-1
+        ).reshape(num_tokens, n_group)
+    group_idx = gating_topk(group_scores, topk_group)[1]  # (num_tokens, top_k_group)
+
+    @T.prim_func(private=True, s_tir=True)
+    def group_limited_mask_scores(
+        var_scores: T.handle, var_group_idx: T.handle, var_output: T.handle
+    ):
+        T.func_attr({"tirx.noalias": True})
+        scores = T.match_buffer(
+            var_scores, (num_tokens, num_routed_experts), dtype=scores_for_choice.dtype
+        )
+        group_idx_tir = T.match_buffer(
+            var_group_idx, (num_tokens, topk_group), dtype=group_idx.dtype
+        )
+        output = T.match_buffer(
+            var_output, (num_tokens, num_routed_experts), dtype=scores_for_choice.dtype
+        )
+        for i, j, k in T.grid(num_tokens, topk_group, group_size):
+            with T.sblock("mask_scores"):
+                vi, vj, vk = T.axis.remap("SSS", [i, j, k])
+                output[vi, group_idx_tir[vi, vj] * group_size + vk] = scores[
+                    vi, group_idx_tir[vi, vj] * group_size + vk
+                ]
+
+    tmp_scores = op.tensor_ir_inplace_op(
+        group_limited_mask_scores,
+        "group_limited_mask_scores",
+        args=[
+            scores_for_choice,
+            group_idx,
+            op.full(
+                scores_for_choice.shape,
+                float(np.finfo("float32").min),
+                dtype=scores_for_choice.dtype,
+            ),
+        ],
+        inplace_indices=[2],
+        out=Tensor.placeholder(scores_for_choice.shape, scores_for_choice.dtype),
+    )
+
+    expert_weights, expert_indices = gating_topk(tmp_scores, top_k)
+    if topk_method == "noaux_tc":
+
+        @T.prim_func(private=True, s_tir=True)
+        def gather_scores(var_scores: T.handle, var_expert_indices: T.handle, var_output: T.handle):
+            T.func_attr({"tirx.noalias": True})
+            scores = T.match_buffer(
+                var_scores,
+                (num_tokens, num_routed_experts),
+                dtype=scores_for_choice.dtype,
+            )
+            expert_indices_tir = T.match_buffer(
+                var_expert_indices, (num_tokens, top_k), dtype=expert_indices.dtype
+            )
+            output = T.match_buffer(var_output, (num_tokens, top_k), dtype=scores_for_choice.dtype)
+            for i, j in T.grid(num_tokens, top_k):
+                with T.sblock("gather_scores"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    output[vi, vj] = scores[vi, expert_indices_tir[vi, vj]]
+
+        expert_weights = op.tensor_ir_op(
+            gather_scores,
+            "gather_scores",
+            args=[scores, expert_indices],
+            out=Tensor.placeholder((num_tokens, top_k), scores_for_choice.dtype),
+        )
+    return expert_weights, expert_indices
+
+
+def moe_cumsum(expert_indices: Tensor, num_local_experts: int) -> Tensor:
+    """An operator that returns the cumsum array in MoE.
+
+    The input `expert_indices` of shape [batch_size, experts_per_tok] indicates the indices of
+    the activated experts for each instance in a batch. This operator first converts it to
+    `expert_mask`, a boolean mask with shape [batch_size, num_local_experts], and then computes
+    cumsum over the transpose-then-flattened array of `expert_mask`.
+
+    A position `(e, b)` in the result `cumsum`, where `e` is the expert id and `b` is the batch id,
+    indicates a shuffling plan that moves the `b`-th instance that ensures the inputs to the `e`-th
+    expert is contiguous.
+
+    Parameters
+    ----------
+    expert_indices : Tensor
+        The topk indices with shape [batch_size, experts_per_tok], int32, where
+        `experts_per_tok` is the number of activated experts.
+
+    num_local_experts : int
+        The number of totally experts.
+
+    Returns
+    -------
+    cumsum: Tensor
+        The cumsum result with shape [num_local_experts * batch_size], int32.
+
+    Example
+    -------
+    Suppose `batch_size` is 4, `experts_per_tok` is 2, the total number of experts is 6, and
+    `expert_indices` is the 2D tensor below:
+
+        [
+            [0, 1],
+            [1, 2],
+            [3, 4],
+            [2, 5],
+        ]
+
+    , then the `expert_mask` is a tensor of shape [batch_size, num_local_experts] below:
+
+        [
+            [1, 1, 0, 0, 0, 0],
+            [0, 1, 1, 0, 0, 0],
+            [0, 0, 0, 1, 1, 0],
+            [0, 0, 1, 0, 0, 1],
+        ]
+
+    . The result cumsum of the transposed `expert_mask` is a flattened version of 2D tensor below:
+
+        [
+            [1, 1, 1, 1],
+            [2, 3, 3, 3],
+            [3, 4, 4, 5],
+            [5, 5, 6, 6],
+            [6, 6, 7, 7],
+            [7, 7, 7, 8],
+        ]
+    """
+    batch_size, experts_per_tok = expert_indices.shape
+    expert_mask = (
+        op.tensor_expr_op(
+            lambda expert_indices: te.compute(
+                (batch_size, num_local_experts),
+                lambda i, j: tirx.expr.Select(
+                    reduce(
+                        tirx.Or,
+                        [expert_indices[i, k] == j for k in range(experts_per_tok)],
+                    ),
+                    true_value=tirx.const(1, "int32"),
+                    false_value=tirx.const(0, "int32"),
+                ),
+            ),
+            "expert_mask",
+            args=[expert_indices],
+        )
+        .permute_dims(1, 0)
+        .reshape(batch_size * num_local_experts)
+    )
+
+    return op.cumsum(expert_mask, axis=0, exclusive=False, dtype="int32")
+
+
+def get_indices(cumsum: Tensor, expert_indices: Tensor) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+    """Returns a 1D tensor of indices that represents the shuffling plan for each instance in a
+    batch, so that the inputs to each experts are contiguous and the indices for reverse permutation
+    (scatter) to the original order.
+
+    If `reverse_indices[i] = (b, j)`, it means the `b`-th instance in the batch should be moved to the
+    `i`-th position in shuffling, and `j` doesn not matter only meaning `expert_indices[b, j]`
+    corresponds to the expert at position `i` in the shuffling plan. We also compute
+    `token_indices[i] = b` so that we can use `relax.op.take` for shuffling.
+
+    Effectively it is equivalent to the following Python code:
+
+    .. code-block:: python
+
+        for b in range(batch_size):
+            for j in range(experts_per_tok):
+                e = expert_indices[b, j]
+                reverse_indices[cumsum[e * batch_size + b] - 1] = b * experts_per_tok + j
+                token_indices[cumsum[e * batch_size + b] - 1
+
+    Parameters
+    ----------
+    cumsum : Tensor
+        A flattened 1D tensor whose original shape is [experts_per_tok, batch_size].
+
+    expert_indices : Tensor
+        The indices of the experts with shape [batch_size, experts_per_tok].
+
+    Returns
+    -------
+    reverse_indices : Tensor
+        The indices for scattering with shape [batch_size * experts_per_tok].
+
+    token_indices : Tensor
+        The indices for shuffling with shape [batch_size * experts_per_tok].
+    """  # noqa: E501
+    TX = 1024
+    batch_size, experts_per_tok = expert_indices.shape
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_cumsum: T.handle,
+        var_expert_indices: T.handle,
+        var_reverse_indices: T.handle,
+        var_token_indices: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        batch_size = T.int32()
+        cumsum_len = T.int32()  # [experts_per_tok * batch_size]
+        cumsum = T.match_buffer(var_cumsum, [cumsum_len], "int32")
+        expert_indices = T.match_buffer(var_expert_indices, [batch_size, experts_per_tok], "int32")
+        reverse_indices = T.match_buffer(
+            var_reverse_indices, [batch_size * experts_per_tok], "int32"
+        )
+        token_indices = T.match_buffer(var_token_indices, [batch_size * experts_per_tok], "int32")
+        for bj_o in T.thread_binding(0, T.ceildiv(batch_size * experts_per_tok, TX), "blockIdx.x"):
+            for bj_i in T.thread_binding(0, TX, "threadIdx.x"):
+                with T.sblock("indices"):
+                    T.reads(expert_indices[:, :], cumsum[:])
+                    T.writes(reverse_indices[:], token_indices[:])
+                    if bj_o * TX + bj_i < batch_size * experts_per_tok:
+                        b: T.int32 = T.floordiv(bj_o * TX + bj_i, experts_per_tok)
+                        j: T.int32 = T.floormod(bj_o * TX + bj_i, experts_per_tok)
+                        e: T.int32 = expert_indices[b, j]
+                        reverse_indices[cumsum[e * batch_size + b] - 1] = b * experts_per_tok + j
+                        token_indices[cumsum[e * batch_size + b] - 1] = b
+
+    return op.tensor_ir_op(
+        _func,
+        "get_indices",
+        args=[cumsum, expert_indices],
+        out=[Tensor.placeholder([batch_size * experts_per_tok], "int32") for _ in range(2)],
+    )
+
+
+def get_indptr(
+    cumsum: Tensor,
+    num_local_experts: int,
+    batch_size: Union[int, tirx.Var],
+    inclusive: bool,
+    out_dtype: str,
+) -> Tensor:
+    """Extract the `indptr` array from MoE cumsum array. The MoE cumsum array is a flattened tensor
+    whose original shape is [num_local_experts, batch_size], and the `indptr` array is a 1D tensor
+    of length `num_local_experts + 1`. The range `[indptr[i], indptr[i + 1])` indicates instances in
+    the batch that corresponds to the `i`-th expert.
+
+    Effectively, this operator is equivalent to the following numpy code:
+
+    .. code-block:: python
+
+        indptr = np.zeros(num_local_experts + 1, dtype=np.int32)
+        indptr[0] = 0
+        for i in range(1, num_local_experts + 1):
+            indptr[i] = cumsum[i * batch_size - 1]
+        return indptr
+
+    Parameters
+    ----------
+    cumsum : Tensor
+        The prefix sum of the sparse array with shape [batch_size * num_local_experts], int32.
+
+    num_local_experts : int
+        The number of experts.
+
+    batch_size : int | tirx.Var
+        The batch size. Note that the batch size here refers to `batch_size * seq_len` in MoE,
+        and we name is `batch_size` for simplicity here only because the two dimensions are fused
+        in Mixtral.
+
+    inclusive : bool
+        Whether to compute inclusive or exclusive prefix sum as the indptr. If `inclusive` is False,
+        the 0-th element of the `indptr` array, which always equals to 0, will be omitted.
+
+    out_dtype : str
+        The output dtype.
+
+    Returns
+    -------
+    indptr : Tensor
+        The `indptr` array with shape [num_local_experts + 1] if `inclusive` is True, otherwise
+        [num_local_experts]. The `indptr` array is of type `out_dtype`.
+    """
+
+    out_shape = [num_local_experts if inclusive else num_local_experts + 1]
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func_exclusive(var_cumsum: T.handle, var_indptr: T.handle, batch_size: T.int64):
+        T.func_attr({"tirx.noalias": True})
+        cumsum = T.match_buffer(var_cumsum, shape=[batch_size * num_local_experts], dtype="int32")
+        indptr = T.match_buffer(var_indptr, shape=out_shape, dtype=out_dtype)
+        for vi in T.serial(0, out_shape[0]):
+            with T.sblock("indptr"):
+                i = T.axis.spatial(out_shape[0], vi)
+                indptr[i] = T.Select(i > 0, cumsum[i * batch_size - 1], T.int32(0))
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func_inclusive(var_cumsum: T.handle, var_indptr: T.handle, batch_size: T.int64):
+        T.func_attr({"tirx.noalias": True})
+        cumsum = T.match_buffer(var_cumsum, shape=[batch_size * num_local_experts], dtype="int32")
+        indptr = T.match_buffer(var_indptr, shape=out_shape, dtype=out_dtype)
+        for vi in T.serial(0, out_shape[0]):
+            with T.sblock("indptr"):
+                i = T.axis.spatial(out_shape[0], vi)
+                indptr[i] = cumsum[(i + 1) * batch_size - 1]
+
+    assert cumsum.ndim == 1
+    return op.tensor_ir_op(
+        _func_inclusive if inclusive else _func_exclusive,
+        "get_expert_instance_indptr",
+        args=[cumsum, batch_size],
+        out=Tensor.placeholder(out_shape, out_dtype),
+    )
+
+
+def scatter_output(x: Tensor, indices: Tensor) -> Tensor:
+    """Scatter the output of MoE experts back to the original positions.
+
+    Parameters
+    ----------
+    x : Tensor
+        The output of MoE experts with shape [batch_size * num_experts_per_tok, hidden_size].
+
+    indices : Tensor
+        The indices of the experts with shape [batch_size * num_experts_per_tok].
+
+    Returns
+    -------
+    out : Tensor
+        The output of MoE experts with shape [batch_size * num_experts_per_tok, hidden_size].
+    """
+    dtype = x.dtype
+    _, hidden_size = x.shape
+
+    @T.prim_func(private=True, s_tir=True)
+    def _func(var_x: T.handle, var_indices: T.handle, var_out: T.handle):
+        T.func_attr({"tirx.noalias": True})
+        indices_len = T.int64()
+        x = T.match_buffer(var_x, [indices_len, hidden_size], dtype)
+        indices = T.match_buffer(var_indices, [indices_len], "int32")
+        out = T.match_buffer(var_out, [indices_len, hidden_size], dtype)
+        for i in T.serial(0, indices_len):
+            for j in T.serial(0, hidden_size):
+                with T.sblock("scatter"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    out[indices[vi], vj] = x[vi, vj]
+
+    return op.tensor_ir_op(
+        _func,
+        "scatter_output",
+        args=[x, indices],
+        out=Tensor.placeholder(x.shape, dtype),
+    )

python/mlc_llm/op/mrope.py

@@ -0,0 +1,442 @@
+"""Utilities for Multimodal Rotary Position Embeddings (MRoPE)."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+from tvm import te, tirx
+from tvm.relax.frontend import nn
+from tvm.relax.frontend.nn import Tensor, op
+
+
+def _rotate_half(x: Tensor) -> Tensor:
+    """Rotate the last dimension of ``x`` by swapping pairs."""
+
+    x1, x2 = op.split(x, 2, axis=-1)
+    return op.concat([op.negative(x2), x1], dim=-1)
+
+
+def _repeat_mrope_section(section: Sequence[int]) -> Tuple[int, ...]:  # noqa: UP006
+    if not section:
+        raise ValueError("mrope_section must not be empty.")
+    if any(s <= 0 for s in section):
+        raise ValueError(f"All mrope_section entries must be positive, got {section}.")
+    return tuple(section) * 2
+
+
+def _split_indices_from_sizes(sizes: Sequence[int]) -> List[int]:  # noqa: UP006
+    indices: List[int] = []  # noqa: UP006
+    running = 0
+    # Drop the final cumulative sum so split() keeps the last chunk.
+    for size in sizes[:-1]:
+        running += size
+        indices.append(running)
+    return indices
+
+
+def _reorder_cos_sin(
+    tensor: Tensor,
+    split_sizes: Sequence[int],
+) -> Tensor:
+    """Reorder cos/sin tensors so the head dimension follows T/H/W repeating sections."""
+
+    if not split_sizes:
+        raise ValueError("split_sizes must not be empty.")
+    split_points = _split_indices_from_sizes(split_sizes)
+    # relax.op.split returns a Python tuple, so we can iterate directly.
+    sections = op.split(tensor, indices_or_sections=split_points, axis=-1)
+    reordered = []
+    for idx, chunk in enumerate(sections):
+        axis_selector = nn.Tensor.from_const(np.array([idx % 3], dtype="int32"))
+        axis_slice = op.take(chunk, axis_selector, axis=0)
+        reordered.append(nn.op.squeeze(axis_slice, 0))
+    return op.concat(reordered, dim=-1)
+
+
+class MultimodalRotaryEmbedding(nn.Module):
+    """Generate cosine/sine tables for multimodal rotary embeddings."""
+
+    def __init__(
+        self,
+        head_dim: int,
+        theta: float,
+        mrope_section: Sequence[int],
+        attention_scaling: float = 1.0,
+    ) -> None:
+        if head_dim % 2 != 0:
+            raise ValueError(f"head_dim must be even for RoPE, got {head_dim}.")
+        self.head_dim = head_dim
+        self.theta = theta
+        self.attention_scaling = attention_scaling
+        self.mrope_section = tuple(mrope_section)
+        self._inv_freq = 1.0 / (
+            theta ** (np.arange(0, head_dim, 2, dtype="float32") / np.float32(head_dim))
+        )
+
+    def forward(self, reference: Tensor, position_ids: Tensor) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+        """Return ``(cos, sin)`` with shape ``(3, batch, seq, head_dim)``."""
+        if len(position_ids.shape) != 3:
+            raise ValueError(
+                "position_ids must be rank-3 with either "
+                "(batch, seq, 3) or (3, batch, seq) layout, "
+                f"got shape {position_ids.shape}."
+            )
+        if isinstance(position_ids.shape[0], int) and position_ids.shape[0] == 3:
+            batch_size, seq_len = position_ids.shape[1], position_ids.shape[2]
+            pos_tensor = op.reshape(position_ids, (3, batch_size, 1, seq_len))
+        elif isinstance(position_ids.shape[-1], int) and position_ids.shape[-1] == 3:
+            batch_size, seq_len = position_ids.shape[0], position_ids.shape[1]
+            permuted_pos = op.permute_dims(position_ids, axes=[2, 0, 1])
+            pos_tensor = op.reshape(permuted_pos, (3, batch_size, 1, seq_len))
+        else:
+            raise ValueError(
+                "position_ids must have exactly one static dimension of size 3, "
+                f"got shape {position_ids.shape}."
+            )
+
+        dtype = reference.dtype
+        inv_freq_tensor = nn.Tensor.from_const(self._inv_freq.reshape(1, 1, -1, 1))
+        inv_freq_tensor = op.broadcast_to(inv_freq_tensor, (3, batch_size, self._inv_freq.size, 1))
+
+        freqs = op.matmul(inv_freq_tensor.astype("float32"), pos_tensor.astype("float32"))
+        freqs = op.permute_dims(freqs, axes=[0, 1, 3, 2])
+        emb = op.concat([freqs, freqs], dim=-1)
+
+        def _apply_trig(func_name: str) -> Tensor:
+            def compute(x: te.Tensor):
+                return te.compute(
+                    x.shape,
+                    lambda *indices: getattr(tirx, func_name)(x[indices]),
+                    name=f"mrope_{func_name}",
+                )
+
+            return op.tensor_expr_op(compute, f"mrope_{func_name}", [emb])
+
+        cos = _apply_trig("cos") * self.attention_scaling
+        sin = _apply_trig("sin") * self.attention_scaling
+        return cos.astype(dtype), sin.astype(dtype)
+
+
+def apply_multimodal_rotary_pos_emb(
+    q: Tensor,
+    k: Tensor,
+    cos: Tensor,
+    sin: Tensor,
+    mrope_section: Sequence[int],
+    unsqueeze_dim: int = 2,
+) -> Tuple[Tensor, Tensor]:  # noqa: UP006
+    """Apply multimodal rotary embedding to query and key tensors."""
+
+    split_sizes = _repeat_mrope_section(mrope_section)
+    reordered_cos = _reorder_cos_sin(cos, split_sizes)
+    reordered_sin = _reorder_cos_sin(sin, split_sizes)
+    cos_term = op.unsqueeze(reordered_cos, dim=unsqueeze_dim)
+    sin_term = op.unsqueeze(reordered_sin, dim=unsqueeze_dim)
+    cos_term = cos_term.astype(q.dtype)
+    sin_term = sin_term.astype(q.dtype)
+    q_embed = op.add(op.multiply(q, cos_term), op.multiply(_rotate_half(q), sin_term))
+    k_embed = op.add(op.multiply(k, cos_term), op.multiply(_rotate_half(k), sin_term))
+    return q_embed, k_embed
+
+
+@dataclass
+class VisionPositionMetadata:
+    """Metadata required to build multimodal position IDs."""
+
+    vision_start_token_id: int
+    image_token_id: int
+    video_token_id: int
+    spatial_merge_size: int
+    tokens_per_second: float
+
+    def merged_hw(self, height: int, width: int) -> Tuple[int, int]:  # noqa: UP006
+        """Return merged height/width after applying ``spatial_merge_size``."""
+
+        if height % self.spatial_merge_size != 0 or width % self.spatial_merge_size != 0:
+            raise ValueError(
+                "Image or video grid is not divisible by spatial_merge_size "
+                f"(got h={height}, w={width}, merge={self.spatial_merge_size})."
+            )
+        return height // self.spatial_merge_size, width // self.spatial_merge_size
+
+
+def _text_chunk(length: int, offset: int) -> np.ndarray:
+    """Create a text-position chunk with a shared scalar offset for T/H/W axes."""
+
+    if length <= 0:
+        return np.zeros((3, 0), dtype=np.int64)
+    seq: np.ndarray = np.arange(length, dtype=np.int64)
+    chunk = np.broadcast_to(seq.reshape(1, -1), (3, length))
+    return chunk + offset
+
+
+def _grid_chunk(
+    grid_t: int,
+    grid_h: int,
+    grid_w: int,
+    offset: int,
+    tokens_per_second: float,
+    second_per_grid: float,
+) -> np.ndarray:
+    if grid_t <= 0 or grid_h <= 0 or grid_w <= 0:
+        raise ValueError(
+            f"Invalid grid shape t={grid_t}, h={grid_h}, w={grid_w} for multimodal positions."
+        )
+    time_axis = (np.arange(grid_t, dtype=np.float32) * second_per_grid * tokens_per_second).astype(
+        np.int64
+    )
+    t_index = np.repeat(time_axis, grid_h * grid_w)
+    h_index = np.tile(np.repeat(np.arange(grid_h, dtype=np.int64), grid_w), grid_t)
+    w_index = np.tile(np.tile(np.arange(grid_w, dtype=np.int64), grid_h), grid_t)
+    stacked = np.stack([t_index, h_index, w_index])
+    return stacked + offset
+
+
+def _find_token_index(tokens: Sequence[int], token_id: int, start: int) -> int:
+    for idx in range(start, len(tokens)):
+        if tokens[idx] == token_id:
+            return idx
+    return len(tokens)
+
+
+def _next_chunk_offset(chunks: Sequence[np.ndarray]) -> int:
+    if not chunks:
+        return 0
+    return int(chunks[-1].max()) + 1
+
+
+def _count_vision_items(
+    token_array: np.ndarray,
+    vision_start_token_id: int,
+    image_token_id: int,
+    video_token_id: int,
+) -> Tuple[int, int]:  # noqa: UP006
+    vision_starts = np.where(token_array == vision_start_token_id)[0]
+    valid_starts = vision_starts[vision_starts + 1 < token_array.shape[0]]
+    following_tokens = token_array[valid_starts + 1]
+    image_count = int(np.sum(following_tokens == image_token_id))
+    video_count = int(np.sum(following_tokens == video_token_id))
+    return image_count, video_count
+
+
+def _next_vision_block(
+    tokens: Sequence[int],
+    start: int,
+    meta: VisionPositionMetadata,
+    has_images: bool,
+    has_videos: bool,
+) -> Tuple[str, int]:  # noqa: UP006
+    sentinel = len(tokens) + 1
+    image_end = _find_token_index(tokens, meta.image_token_id, start) if has_images else sentinel
+    video_end = _find_token_index(tokens, meta.video_token_id, start) if has_videos else sentinel
+    if image_end < video_end:
+        return "image", image_end
+    return "video", video_end
+
+
+def _load_grid_for_block(
+    block_kind: str,
+    image_grid_thw: Optional[np.ndarray],  # noqa: UP045
+    video_grid_thw: Optional[np.ndarray],  # noqa: UP045
+    second_per_grid_ts: Optional[np.ndarray],  # noqa: UP045
+    image_index: int,
+    video_index: int,
+) -> Tuple[int, int, int, float, int, int]:  # noqa: UP006
+    if block_kind == "image":
+        if image_grid_thw is None:
+            raise ValueError("Image grids are required for sequences with image tokens.")
+        grid_t, grid_h, grid_w = image_grid_thw[image_index]
+        return int(grid_t), int(grid_h), int(grid_w), 0.0, image_index + 1, video_index
+
+    if video_grid_thw is None:
+        raise ValueError("Video grids are required for sequences with video tokens.")
+    grid_t, grid_h, grid_w = video_grid_thw[video_index]
+    second_per_grid = (
+        float(second_per_grid_ts[video_index]) if second_per_grid_ts is not None else 1.0
+    )
+    return int(grid_t), int(grid_h), int(grid_w), second_per_grid, image_index, video_index + 1
+
+
+def _build_sequence_position_ids(
+    input_tokens: Sequence[int],
+    meta: VisionPositionMetadata,
+    image_grid_thw: Optional[np.ndarray],  # noqa: UP045
+    video_grid_thw: Optional[np.ndarray],  # noqa: UP045
+    second_per_grid_ts: Optional[np.ndarray],  # noqa: UP045
+    image_index: int,
+    video_index: int,
+) -> Tuple[np.ndarray, int, int, int]:  # noqa: UP006
+    token_array = np.asarray(input_tokens, dtype=np.int64)
+    image_count, video_count = _count_vision_items(
+        token_array,
+        vision_start_token_id=meta.vision_start_token_id,
+        image_token_id=meta.image_token_id,
+        video_token_id=meta.video_token_id,
+    )
+    if image_count > 0 and image_grid_thw is None:
+        raise ValueError("Image grids are required for sequences with image tokens.")
+    if video_count > 0 and video_grid_thw is None:
+        raise ValueError("Video grids are required for sequences with video tokens.")
+
+    llm_pos_ids_list: List[np.ndarray] = []  # noqa: UP006
+    start = 0
+    remain_images = image_count
+    remain_videos = video_count
+    for _ in range(image_count + video_count):
+        block_kind, block_end = _next_vision_block(
+            tokens=input_tokens,
+            start=start,
+            meta=meta,
+            has_images=remain_images > 0,
+            has_videos=remain_videos > 0,
+        )
+        (
+            grid_t,
+            grid_h,
+            grid_w,
+            second_per_grid,
+            image_index,
+            video_index,
+        ) = _load_grid_for_block(
+            block_kind=block_kind,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            image_index=image_index,
+            video_index=video_index,
+        )
+        if block_kind == "image":
+            remain_images -= 1
+        else:
+            remain_videos -= 1
+
+        llm_grid_h, llm_grid_w = meta.merged_hw(grid_h, grid_w)
+        text_len = block_end - start
+        text_offset = _next_chunk_offset(llm_pos_ids_list)
+        llm_pos_ids_list.append(_text_chunk(text_len, text_offset))
+        grid_offset = text_offset + text_len
+        llm_pos_ids_list.append(
+            _grid_chunk(
+                grid_t=grid_t,
+                grid_h=llm_grid_h,
+                grid_w=llm_grid_w,
+                offset=grid_offset,
+                tokens_per_second=meta.tokens_per_second,
+                second_per_grid=second_per_grid,
+            )
+        )
+        start = block_end + grid_t * llm_grid_h * llm_grid_w
+
+    if start < len(input_tokens):
+        tail_len = len(input_tokens) - start
+        tail_offset = _next_chunk_offset(llm_pos_ids_list)
+        llm_pos_ids_list.append(_text_chunk(tail_len, tail_offset))
+
+    if not llm_pos_ids_list:
+        empty_positions: np.ndarray = np.zeros((3, 0), dtype=np.int64)
+        return empty_positions, 0, image_index, video_index
+    llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
+    delta = int(llm_positions.max()) + 1 - len(input_tokens)
+    return llm_positions, delta, image_index, video_index
+
+
+def _text_only_position_ids(
+    input_ids: np.ndarray,
+    attention_mask: Optional[np.ndarray],  # noqa: UP045
+) -> Tuple[np.ndarray, np.ndarray]:  # noqa: UP006
+    batch, seq_len = input_ids.shape
+    if attention_mask is None:
+        base: np.ndarray = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1)
+        tiled = np.broadcast_to(base, (3, batch, seq_len))
+        return tiled, np.zeros((batch, 1), dtype=np.int64)
+
+    position = attention_mask.cumsum(axis=-1) - 1
+    position = np.where(attention_mask == 0, 1, position)
+    position = np.expand_dims(position, axis=0).repeat(3, axis=0)
+    max_pos = position.max(axis=0, keepdims=False).max(axis=-1, keepdims=True)
+    delta = (max_pos + 1 - seq_len).astype(np.int64)
+    return position.astype(np.int64), delta
+
+
+def get_mrope_position_ids(
+    input_ids: np.ndarray,
+    meta: VisionPositionMetadata,
+    attention_mask: Optional[np.ndarray] = None,  # noqa: UP045
+    image_grid_thw: Optional[np.ndarray] = None,  # noqa: UP045
+    video_grid_thw: Optional[np.ndarray] = None,  # noqa: UP045
+    second_per_grid_ts: Optional[np.ndarray] = None,  # noqa: UP045
+) -> Tuple[np.ndarray, np.ndarray]:  # noqa: UP006
+    """Generate 3D position IDs and deltas following Hugging Face Qwen2.5-VL."""
+
+    input_ids = np.asarray(input_ids, dtype=np.int64)
+    batch, seq_len = input_ids.shape
+    position_ids = np.ones((3, batch, seq_len), dtype=np.int64)
+
+    attention = None
+    if attention_mask is not None:
+        attention_mask = np.asarray(attention_mask, dtype=np.int64)
+        if attention_mask.shape != input_ids.shape:
+            raise ValueError(
+                "attention_mask shape must match input_ids shape: "
+                f"{attention_mask.shape} vs {input_ids.shape}"
+            )
+        attention = attention_mask.astype(bool)
+
+    image_grid_thw = None if image_grid_thw is None else np.asarray(image_grid_thw, dtype=np.int64)
+    video_grid_thw = None if video_grid_thw is None else np.asarray(video_grid_thw, dtype=np.int64)
+    if second_per_grid_ts is not None:
+        second_per_grid_ts = np.asarray(second_per_grid_ts, dtype=np.float32)
+
+    contains_image_tokens = bool(np.any(input_ids == meta.image_token_id))
+    contains_video_tokens = bool(np.any(input_ids == meta.video_token_id))
+    if contains_image_tokens and image_grid_thw is None:
+        raise ValueError("image_grid_thw must be provided when image tokens exist in input_ids.")
+    if contains_video_tokens and video_grid_thw is None:
+        raise ValueError("video_grid_thw must be provided when video tokens exist in input_ids.")
+    if (
+        second_per_grid_ts is not None
+        and video_grid_thw is not None
+        and second_per_grid_ts.shape[0] != video_grid_thw.shape[0]
+    ):
+        raise ValueError(
+            "second_per_grid_ts length must match number of video grids "
+            f"({second_per_grid_ts.shape[0]} vs {video_grid_thw.shape[0]})."
+        )
+
+    if not (contains_image_tokens or contains_video_tokens):
+        return _text_only_position_ids(input_ids, attention_mask)
+
+    image_index = 0
+    video_index = 0
+    deltas: List[int] = []  # noqa: UP006
+
+    for batch_idx in range(batch):
+        tokens = input_ids[batch_idx]
+        if attention is not None:
+            tokens = tokens[attention[batch_idx]]
+        token_values = np.asarray(tokens, dtype=np.int64).tolist()
+        input_tokens: List[int] = [int(token) for token in token_values]  # noqa: UP006
+        if not input_tokens:
+            deltas.append(0)
+            continue
+
+        llm_positions, delta, image_index, video_index = _build_sequence_position_ids(
+            input_tokens=input_tokens,
+            meta=meta,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            image_index=image_index,
+            video_index=video_index,
+        )
+        if attention is not None:
+            position_ids[:, batch_idx, attention[batch_idx]] = llm_positions
+        else:
+            position_ids[:, batch_idx, :] = llm_positions
+        deltas.append(delta)
+
+    delta_array = np.asarray(deltas, dtype=np.int64).reshape(batch, 1)
+    return position_ids, delta_array

python/mlc_llm/op/pipeline_parallel.py

@@ -0,0 +1,33 @@
+"""Operators for pipeline parallelism."""
+
+from typing import List  # noqa: UP035
+
+from tvm import relax
+from tvm.relax.frontend.nn import Tensor, op
+
+
+def pipeline_stage_boundary(*tensors: Tensor) -> List[Tensor]:  # noqa: UP006
+    """Pipeline parallelism stage boundary mark operator in MLC.
+
+    Parameters
+    ----------
+    tensors : Tensor
+        The tensors to be passed to the next stage.
+
+    Returns
+    -------
+    tensors : List[Tensor]
+        The list of input tensors passed to the next stage.
+    """
+    return op.wrap_nested(
+        relax.call_pure_packed(
+            "mlc.pipeline_parallel_stage_boundary",
+            *[tensor._expr for tensor in tensors],
+            ty_args=(
+                tensors[0]._expr.ty
+                if len(tensors) == 1
+                else relax.TupleType([tensor._expr.ty for tensor in tensors])
+            ),
+        ),
+        name="pipeline_stage_boundary",
+    )

python/mlc_llm/op/top_p_pivot.py

@@ -0,0 +1,335 @@
+"""Operators for choosing the pivot to cut-off top-p percentile"""
+
+import tvm
+from tvm.script import tirx as T
+
+from mlc_llm.support.max_thread_check import get_max_num_threads_per_block
+
+# mypy: disable-error-code="attr-defined,valid-type,name-defined"
+
+
+def top_p_pivot(pN, target: tvm.target.Target):
+    """Top-p pivot function. This function finds the pivot to cut-off top-p percentile.
+
+    A valide pivot should satisfy the following conditions:
+    - lsum >= top_p
+    - top_p > lsum - cmin * lmin
+    where lsum is the sum of elements that are larger or equal to the pivot,
+    lmin is the minimum elements that is larger or equal to the pivot,
+    cmin is the count of elements that are equal to lmin,
+
+    Parameters
+    ----------
+    prob:
+        The probability vector
+
+    top_p_arr:
+        The top-p threshold
+
+    init_pivots:
+        The initial pivot candidates
+
+    final_pivot:
+        The final pivot to cut-off top-p percentile
+
+    final_lsum:
+        The final sum of the values after top-p filtering.
+    """
+    TX = 1024
+    K = 32
+    eps_LR = 1e-7
+
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    TX = min(TX, max_num_threads_per_block)
+
+    def _var(dtype="int32"):
+        return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+    def valid(lsum, lmin, cmin, top_p):
+        return tvm.tirx.all(lsum >= top_p, top_p > lsum - cmin * lmin)
+
+    # fmt: off
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_prob: T.handle,
+        var_top_p_arr: T.handle,
+        var_init_pivots: T.handle,
+        var_final_pivot: T.handle,
+        var_final_lsum: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        B = T.int32()
+        N = T.int32()
+        prob = T.match_buffer(var_prob, (B, N,), "float32")
+        top_p_arr = T.match_buffer(var_top_p_arr, (B,), dtype="float32")
+        init_pivots = T.match_buffer(var_init_pivots, (B, pN), "float32")
+        final_pivot = T.match_buffer(var_final_pivot, (B,), "float32")
+        final_lsum = T.match_buffer(var_final_lsum, (B,), "float32")
+
+        with T.sblock("kernel"):
+            pivot = T.sblock_alloc_buffer((pN,), "float32", scope="local")
+            top_p = _var("float32")
+
+            L = T.sblock_alloc_buffer((1,), "float32", scope="shared")
+            R = T.sblock_alloc_buffer((1,), "float32", scope="shared")
+            L_local = _var("float32")
+            R_local = _var("float32")
+
+            q = _var("float32")
+            lsum = T.sblock_alloc_buffer((pN,), "float32", scope="local")
+            lmin_broadcast = T.sblock_alloc_buffer((1), "float32", scope="shared")
+            lmin_broadcast_local = _var("float32")
+            lmin = T.sblock_alloc_buffer((pN,), "float32", scope="local")
+            cmin = T.sblock_alloc_buffer((pN,), "int32", scope="local")
+            total_sum = _var("float32")
+
+            it = _var("int32")
+            es_local = _var("bool")
+            es = T.sblock_alloc_buffer((1,), "bool", scope="shared")
+            find_pivot_local = _var("bool")
+            find_pivot = T.sblock_alloc_buffer((1,), "bool", scope="shared")
+
+            total_sum_reduce = _var("float32")
+            lsum_reduce = _var("float32")
+            lmin_reduce = _var("float32")
+            cmin_reduce = _var("int32")
+
+            for _bx in T.thread_binding(0, B, thread="blockIdx.x"):
+                for _tx in T.thread_binding(0, TX, thread="threadIdx.x"):
+                    with T.sblock("CTA"):
+                        b, tx = T.axis.remap("SS", [_bx, _tx])
+
+                        top_p[0] = top_p_arr[b]
+
+                        if tx == 0:
+                            # leader thread initializes L, R
+                            L[0] = 1.0 - top_p[0]
+                            R[0] = eps_LR
+                            find_pivot[0] = False
+                        T.tvm_storage_sync("shared")
+
+                        L_local[0] = L[0]
+                        R_local[0] = R[0]
+                        for i in T.unroll(0, pN):
+                            # pivots are in descending order
+                            pivot[i] = init_pivots[b, i]
+                        find_pivot_local[0] = False
+                        if L_local[0] - R_local[0] <= eps_LR:
+                            # When the initial value is too small, set the result directly.
+                            if tx == 0:
+                                final_lsum[b] = 1.0
+                                final_pivot[b] = 0.0
+                            find_pivot_local[0] = True
+
+                        while T.tvm_thread_invariant(
+                            L_local[0] - R_local[0] > eps_LR
+                            and T.Not(find_pivot_local[0])
+                        ):
+                            # sync before each iteration
+                            T.tvm_storage_sync("shared")
+
+                            ### get lsum, lmin, total_sum
+                            for pidx in T.unroll(0, pN):
+                                lsum[pidx] = 0.0
+                                lmin[pidx] = T.max_value("float32")
+                                cmin[pidx] = 0
+                            total_sum[0] = 0.0
+                            it[0] = 0
+                            es_local[0] = False
+                            while it[0] < T.ceildiv(N, TX) and T.Not(es_local[0]):
+                                idx = T.meta_var(it[0] * TX + tx)
+                                q[0] = T.if_then_else(idx < N, prob[b, idx], 0.0)
+                                total_sum[0] += q[0]
+                                for pidx in T.unroll(0, pN):
+                                    if q[0] >= pivot[pidx]:
+                                        lsum[pidx] += q[0]
+                                        if lmin[pidx] > q[0]:
+                                            lmin[pidx] = q[0]
+                                            cmin[pidx] = 1
+                                        elif lmin[pidx] == q[0]:
+                                            cmin[pidx] += 1
+                                it[0] += 1
+
+                                # early stop every K iterations
+                                if it[0] % K == 0:
+                                    # reduce total_sum over tx
+                                    # T.tvm_storage_sync("shared")
+                                    with T.sblock("block_cross_thread"):
+                                        T.reads(total_sum[0])
+                                        T.writes(total_sum_reduce[0])
+                                        T.attr(
+                                            T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                                            "reduce_scope",
+                                            T.int32(0),
+                                        )
+                                        T.tvm_thread_allreduce(T.uint32(1), total_sum[0], True, total_sum_reduce[0], tx, dtype="void")  # noqa: E501
+                                    # T.tvm_storage_sync("shared")
+
+                                    if tx == 0:
+                                        # leader thread checks if we can stop early
+                                        es[0] = 1 - total_sum_reduce[0] < pivot[pN - 1]
+                                    T.tvm_storage_sync("shared")
+                                    es_local[0] = es[0]
+
+                            T.tvm_storage_sync("shared")
+
+                            # reduce lsum, lmin, cmin, over tx
+                            for pidx in T.serial(0, pN):
+                                # reduce lsum over tx for pivot[j]
+                                with T.sblock("block_cross_thread"):
+                                    T.reads(lsum[pidx])
+                                    T.writes(lsum_reduce[0])
+                                    T.attr(
+                                        T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                                        "reduce_scope",
+                                        T.int32(0),
+                                    )
+                                    T.tvm_thread_allreduce(T.uint32(1), lsum[pidx], True, lsum_reduce[0], tx, dtype="void")  # noqa: E501
+
+                                # reduce lmin over tx for pivot[j]
+                                with T.sblock("block_cross_thread"):
+                                    T.reads(lmin[pidx])
+                                    T.writes(lmin_reduce[0])
+                                    T.attr(
+                                        T.comm_reducer(lambda x0, y0: T.min(x0, y0), [T.float32(0)]),  # noqa: E501
+                                        "reduce_scope",
+                                        T.int32(0),
+                                    )
+                                    T.tvm_thread_allreduce(T.uint32(1), lmin[pidx], True, lmin_reduce[0], tx, dtype="void")  # noqa: E501
+
+                                if tx == 0:
+                                    # broadcast lmin to all threads
+                                    lmin_broadcast[0] = lmin_reduce[0]
+                                T.tvm_storage_sync("shared")
+                                lmin_broadcast_local[0] = lmin_broadcast[0]
+                                if lmin[pidx] > lmin_broadcast_local[0]:
+                                    cmin[pidx] = 0
+                                if tx == 0:
+                                    # only the leader thread updates lsum, lmin
+                                    lsum[pidx] = lsum_reduce[0]
+                                    lmin[pidx] = lmin_reduce[0]
+
+                                # reduce cmin over tx for pivot[j]
+                                with T.sblock("block_cross_thread"):
+                                    T.reads(cmin[pidx])
+                                    T.writes(cmin_reduce[0])
+                                    T.attr(
+                                        T.comm_reducer(lambda x0, y0: x0 + y0, [T.int32(0)]),
+                                        "reduce_scope",
+                                        T.int32(0),
+                                    )
+                                    T.tvm_thread_allreduce(T.uint32(1), cmin[pidx], True, cmin_reduce[0], tx, dtype="void")  # noqa: E501
+
+                                if tx == 0:
+                                    # only the leader thread updates cmin
+                                    cmin[pidx] = cmin_reduce[0]
+
+                            T.tvm_storage_sync("shared")
+
+                            if tx == 0:
+                                # leader thread checks if we have found the pivot, or updates L, R
+                                it[0] = 0
+                                while it[0] < pN and T.Not(find_pivot_local[0]):
+                                    pidx = T.meta_var(it[0])
+                                    if valid(lsum[pidx], lmin[pidx], cmin[pidx], top_p[0]):
+                                        find_pivot[0] = True
+                                        find_pivot_local[0] = True
+                                        # write back the pivot and lsum
+                                        final_pivot[b] = pivot[pidx]
+                                        final_lsum[b] = lsum[pidx]
+                                    elif lsum[pidx] - lmin[pidx] * cmin[pidx] >= top_p[0]:
+                                        R[0] = pivot[pidx]
+                                        final_lsum[b] = lsum[pidx]
+                                    elif lsum[pidx] < top_p[0]:
+                                        L[0] = pivot[pidx]
+                                    it[0] += 1
+
+                            T.tvm_storage_sync("shared")
+
+                            L_local[0] = L[0]
+                            R_local[0] = R[0]
+                            find_pivot_local[0] = find_pivot[0]
+                            # new pivots for next iteration
+                            # uniform spacing between L and R
+                            for pidx in T.unroll(0, pN):
+                                pivot[pidx] = L[0] - (pidx + 1) * (L_local[0] - R_local[0]) / (pN + 1)  # noqa: E501
+
+                        if tx == 0:
+                            # leader thread writes back the pivot
+                            if T.Not(find_pivot_local[0]):
+                                final_pivot[b] = R_local[0]
+                                if R_local[0] == eps_LR:
+                                    final_lsum[b] = lsum[pN - 1]
+    # fmt: on
+
+    return _func
+
+
+def top_p_renorm(target: tvm.target.Target = None):
+    """Top-p renormalization function. This function renormalizes the probability vector.
+
+    Given the pivot, the probability vector is renormalized as follows:
+    - if prob >= pivot, renorm_prob = prob / lsum
+    - otherwise, renorm_prob = 0
+
+    Parameters
+    ----------
+    prob:
+        The probability vector
+
+    final_pivot:
+        The final pivot to cut-off top-p percentile
+
+    final_lsum:
+        The sum of elements that are larger or equal to the pivot
+
+    renorm_prob:
+        The renormalized probability vector
+    """
+    TX = 1024
+    CTA_COUNT = 512
+
+    if target:
+        max_num_threads_per_block = get_max_num_threads_per_block(target)
+        TX = min(TX, max_num_threads_per_block)
+
+    def _var(dtype="int32"):
+        return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+    # fmt: off
+    @T.prim_func(private=True, s_tir=True)
+    def _func(
+        var_prob: T.handle,
+        var_final_pivot: T.handle,
+        var_final_lsum: T.handle,
+        var_renorm_prob: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True})
+        B = T.int32()
+        N = T.int32()
+        prob = T.match_buffer(var_prob, (B, N,), "float32")
+        final_pivot = T.match_buffer(var_final_pivot, (B,), "float32")
+        final_lsum = T.match_buffer(var_final_lsum, (B,), "float32")
+        renorm_prob = T.match_buffer(var_renorm_prob, (B, N,), "float32")
+
+        with T.sblock("kernel"):
+            pivot = _var("float32")
+            lsum = _var("float32")
+            BX = T.meta_var(T.ceildiv(CTA_COUNT, B))
+
+            for _by in T.thread_binding(0, B, thread="blockIdx.y"):
+                for _bx in T.thread_binding(0, BX, thread="blockIdx.x"):
+                    for _tx in T.thread_binding(0, TX, thread="threadIdx.x"):
+                        with T.sblock("CTA"):
+                            by, bx, tx = T.axis.remap("SSS", [_by, _bx, _tx])
+
+                            pivot[0] = final_pivot[by]
+                            lsum[0] = final_lsum[by]
+
+                            for i in T.serial(T.ceildiv(N, BX * TX)):
+                                idx = T.meta_var(i * BX * TX + bx * TX + tx)
+                                if idx < N:
+                                    renorm_prob[by, idx] = T.if_then_else(prob[by, idx] >= pivot[0], prob[by, idx] / lsum[0], 0.0)  # noqa: E501
+    # fmt: on
+
+    return _func

python/mlc_llm/op/triton.py

@@ -0,0 +1,762 @@
+"""Operators enabled by external modules."""
+
+from typing import List, Literal, Tuple  # noqa: UP035
+
+import tvm
+from tvm.relax.frontend import nn
+from tvm.script import ir as I
+from tvm.script import tirx as T
+
+try:
+    import triton
+    import triton.language as tl
+except ImportError:
+    triton = None
+    tl = None
+
+
+# We use a wrapper function to avoid type annotation issue of "tl.constexpr" when
+# triton is not installed.
+def _get_triton_w8a8_block_fp8_gemm():
+    # Triton kernel adapted from SGLang project
+    # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501
+    def _triton_w8a8_block_fp8_gemm(
+        # Pointers to inputs and output
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        # Shape for matmul
+        M,
+        N: tl.constexpr,
+        K: tl.constexpr,
+        # Stride for inputs and output
+        stride_am: tl.constexpr,
+        stride_ak: tl.constexpr,
+        stride_bk: tl.constexpr,
+        stride_bn: tl.constexpr,
+        stride_cm: tl.constexpr,
+        stride_cn: tl.constexpr,
+        stride_As_m: tl.constexpr,
+        stride_As_k: tl.constexpr,
+        stride_Bs_k: tl.constexpr,
+        stride_Bs_n: tl.constexpr,
+        # Block size for block-wise quantization
+        group_n: tl.constexpr,
+        group_k: tl.constexpr,
+        # Meta-parameters
+        BLOCK_SIZE_M: tl.constexpr,
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_K: tl.constexpr,
+        GROUP_SIZE_M: tl.constexpr,
+    ):
+        """Triton-accelerated function used to perform linear operations (dot
+        product) on input tensors `A` and `B` with block-wise quantization,
+        and store the result in output tensor `C`.
+        """
+
+        pid = tl.program_id(axis=0).to(tl.int64)
+        num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+        num_pid_in_group = GROUP_SIZE_M * num_pid_n
+        group_id = pid // num_pid_in_group
+        first_pid_m = group_id * GROUP_SIZE_M
+        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+        pid_m = first_pid_m + (pid % group_size_m)
+        pid_n = (pid % num_pid_in_group) // group_size_m
+
+        offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+        offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+        offs_k = tl.arange(0, BLOCK_SIZE_K)
+        a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+        b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+
+        As_ptrs = As + offs_am * stride_As_m
+        offs_bsn = offs_bn // group_n
+        Bs_ptrs = Bs + offs_bsn * stride_Bs_n
+
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+            a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+            k_start = k * BLOCK_SIZE_K
+            offs_ks = k_start // group_k
+            a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+            b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
+
+            accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
+            a_ptrs += BLOCK_SIZE_K * stride_ak
+            b_ptrs += BLOCK_SIZE_K * stride_bk
+
+        if C.dtype.element_ty == tl.bfloat16:
+            c = accumulator.to(tl.bfloat16)
+        elif C.dtype.element_ty == tl.float16:
+            c = accumulator.to(tl.float16)
+        else:
+            c = accumulator.to(tl.float32)
+
+        offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+        tl.store(c_ptrs, c, mask=c_mask)
+
+    return _triton_w8a8_block_fp8_gemm
+
+
+# We use a wrapper function to avoid type annotation issue of "tl.constexpr" when
+# triton is not installed.
+def _get_triton_w8a8_block_fp8_group_gemm():
+    # Triton kernel adapted from SGLang project
+    # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501
+    def _triton_w8a8_block_fp8_group_gemm(
+        # Pointers to matrices
+        a_ptr,
+        b_ptr,
+        c_ptr,
+        a_scale_ptr,
+        b_scale_ptr,
+        expert_ids_ptr,
+        indptr_ptr,
+        # Matrix dimensions
+        EM,
+        N: tl.constexpr,
+        K: tl.constexpr,
+        num_experts: tl.constexpr,
+        # The stride variables represent how much to increase the ptr by when
+        # moving by 1 element in a particular dimension. E.g. `stride_am` is
+        # how much to increase `a_ptr` by to get the element one row down
+        # (A has M rows).
+        stride_am: tl.constexpr,
+        stride_ak: tl.constexpr,
+        stride_be: tl.constexpr,
+        stride_bk: tl.constexpr,
+        stride_bn: tl.constexpr,
+        stride_cm: tl.constexpr,
+        stride_cn: tl.constexpr,
+        stride_asm: tl.constexpr,
+        stride_ask: tl.constexpr,
+        stride_bse: tl.constexpr,
+        stride_bsk: tl.constexpr,
+        stride_bsn: tl.constexpr,
+        # Block size for block-wise quantization
+        group_n: tl.constexpr,
+        group_k: tl.constexpr,
+        # Meta-parameters
+        BLOCK_SIZE_M: tl.constexpr,
+        BLOCK_SIZE_N: tl.constexpr,
+        BLOCK_SIZE_K: tl.constexpr,
+        GROUP_SIZE_M: tl.constexpr,
+        even_Ks: tl.constexpr,
+    ):
+        """
+        Implements the fused computation for a Mixture of Experts (MOE) using
+        token and expert matrices.
+
+        Key Parameters:
+        - A: The input tensor representing tokens with shape (*, K), where '*' can
+            be any shape representing batches and K is the feature dimension of
+            each token.
+        - B: The stacked MOE weight tensor with shape (E, N, K), where E is
+            the number of experts, K is the input feature dimension, and N is
+            the output feature dimension.
+        - C: The output cache tensor with shape (*, N), where '*' means the
+            same shape as the input tensor A, and N is the output feature dimension.
+        - expert_ids: A tensor containing the indices of the expert for each
+            block. It determines which expert matrix from B should be used for
+            each block in A.
+        This kernel performs the multiplication of a token by its corresponding
+        expert matrix as determined by `expert_ids`.
+        """
+        # -----------------------------------------------------------
+        # Map program ids `pid` to the block of C it should compute.
+        # This is done in a grouped ordering to promote L2 data reuse.
+        pid = tl.program_id(axis=0).to(tl.int64)
+        num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_experts
+        num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+        num_pid_in_group = GROUP_SIZE_M * num_pid_n
+        group_id = pid // num_pid_in_group
+        first_pid_m = group_id * GROUP_SIZE_M
+        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+        pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+        pid_n = (pid % num_pid_in_group) // group_size_m
+
+        # ----------------------------------------------------------
+        # Create pointers for the first blocks of A and B.
+        # We will advance this pointer as we move in the K direction
+        # and accumulate
+        # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+        # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+        expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+        if expert_id == -1:
+            return
+
+        token_begin = tl.load(indptr_ptr + expert_id)
+        token_end = tl.load(indptr_ptr + expert_id + 1)
+        start_pid_m = tl.cdiv(token_begin, BLOCK_SIZE_M) + expert_id
+        offs_token_id = (
+            token_begin + (pid_m - start_pid_m) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        )
+        token_mask = offs_token_id < token_end
+
+        offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+        offs_k = tl.arange(0, BLOCK_SIZE_K)
+        a_ptrs = a_ptr + offs_token_id[:, None] * stride_am + offs_k[None, :] * stride_ak
+
+        b_ptrs = (
+            b_ptr
+            + expert_id * stride_be
+            + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+        )
+
+        a_scale_ptrs = a_scale_ptr + offs_token_id * stride_asm
+        offs_bsn = offs_bn // group_n
+        b_scale_ptrs = b_scale_ptr + expert_id * stride_bse + offs_bsn * stride_bsn
+
+        # -----------------------------------------------------------
+        # Iterate to compute a block of the C matrix.
+        # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+        # of fp32 values for higher accuracy.
+        # `accumulator` will be converted back to fp16 after the loop.
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+            # Load the next block of A and B, generate a mask by checking the
+            # K dimension.
+            if even_Ks:
+                a = tl.load(
+                    a_ptrs,
+                    mask=token_mask[:, None],
+                    other=0.0,
+                )
+                b = tl.load(b_ptrs)
+            else:
+                a = tl.load(
+                    a_ptrs,
+                    mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+                    other=0.0,
+                )
+                b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+            # We accumulate along the K dimension.
+            k_start = k * BLOCK_SIZE_K
+            offs_ks = k_start // group_k
+            a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0)
+            b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
+
+            accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
+            # Advance the ptrs to the next K block.
+            a_ptrs += BLOCK_SIZE_K * stride_ak
+            b_ptrs += BLOCK_SIZE_K * stride_bk
+
+        if c_ptr.dtype.element_ty == tl.bfloat16:
+            accumulator = accumulator.to(tl.bfloat16)
+        elif c_ptr.dtype.element_ty == tl.float16:
+            accumulator = accumulator.to(tl.float16)
+        else:
+            accumulator = accumulator.to(tl.float32)
+
+        # -----------------------------------------------------------
+        # Write back the block of the output
+        offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        c_ptrs = c_ptr + stride_cm * offs_token_id[:, None] + stride_cn * offs_cn[None, :]
+        c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+        tl.store(c_ptrs, accumulator, mask=c_mask)
+
+    return _triton_w8a8_block_fp8_group_gemm
+
+
+def get_tir_w8a8_block_fp8_matmul(
+    N: int,
+    K: int,
+    block_n: int,
+    block_k: int,
+    in_dtype: Literal["float8_e4m3fn"],
+    out_dtype: Literal["float16", "bfloat16"],
+    BLOCK_SIZE_M: int,
+    BLOCK_SIZE_N: int,
+    BLOCK_SIZE_K: int,
+    GROUP_SIZE_M: int,
+    num_warps: int,
+    num_stages: int,
+    extern_mods: List[tvm.runtime.Module],  # noqa: UP006
+):
+    """Get the TIR function for the w8a8_block_fp8_matmul kernel."""
+    # NOTE: adding the type annotation of " -> Tuple[Optional[tvm.tirx.PrimFunc], str]"
+    # will cause the failure of the type resolution in mypy.
+    if triton is None:
+        raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.")
+
+    name_suffix = f"_N{N}_K{K}_block_n{block_n}_block_k{block_k}_in{in_dtype}_out{out_dtype}"
+    kernel_name = f"triton_w8a8_block_fp8_gemm{name_suffix}"
+    tir_name = f"tir_w8a8_block_fp8_matmul{name_suffix}"
+    for ext_mod in extern_mods:
+        if ext_mod.implements_function(kernel_name):
+            return [None, tir_name]
+
+    triton_kernel = _get_triton_w8a8_block_fp8_gemm()
+    triton_kernel.__name__ = kernel_name
+
+    @I.ir_module
+    class BlockFP8Matmul:
+        @T.prim_func(private=True, s_tir=True)
+        def tir_w8a8_block_fp8_matmul(
+            var_A: T.handle,
+            var_B: T.handle,
+            var_As: T.handle,
+            var_Bs: T.handle,
+            var_C: T.handle,
+        ):
+            T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1})
+            M = T.int32()
+            A = T.match_buffer(var_A, (M, K), dtype=in_dtype)
+            B = T.match_buffer(var_B, (N, K), dtype=in_dtype)
+            As = T.match_buffer(var_As, (M, (K + block_k - 1) // block_k), "float32")
+            Bs = T.match_buffer(
+                var_Bs,
+                ((N + block_n - 1) // block_n, (K + block_k - 1) // block_k),
+                "float32",
+            )
+            C = T.match_buffer(var_C, (M, N), dtype=out_dtype)
+            with T.sblock("root"):
+                T.reads(
+                    A[0:M, 0:K],
+                    B[0:N, 0:K],
+                    As[0:M, 0 : (K + block_k - 1) // block_k],
+                    Bs[
+                        0 : (N + block_n - 1) // block_n,
+                        0 : (K + block_k - 1) // block_k,
+                    ],
+                )
+                T.writes(C[0:M, 0:N])
+                T.call_kernel(
+                    triton.jit(triton_kernel),
+                    (T.ceildiv(M, BLOCK_SIZE_M) * T.ceildiv(N, BLOCK_SIZE_N),),
+                    A.data,
+                    B.data,
+                    C.data,
+                    As.data,
+                    Bs.data,
+                    M,
+                    N,
+                    K,
+                    K,  # stride_am
+                    1,  # stride_ak
+                    1,  # stride_bk
+                    K,  # stride_bn
+                    N,  # stride_cm
+                    1,  # stride_cn
+                    (K + block_k - 1) // block_k,  # stride_As_m
+                    1,  # stride_As_k
+                    1,  # stride_Bs_k
+                    (K + block_k - 1) // block_k,  # stride_Bs_n
+                    block_n,
+                    block_k,
+                    BLOCK_SIZE_M,
+                    BLOCK_SIZE_N,
+                    BLOCK_SIZE_K,
+                    GROUP_SIZE_M,
+                    num_warps=num_warps,
+                    num_stages=num_stages,
+                )
+
+    new_ext_mods = BlockFP8Matmul.attrs["external_mods"]
+    assert len(new_ext_mods) == 1
+    extern_mods.append(new_ext_mods[0])
+    return BlockFP8Matmul["tir_w8a8_block_fp8_matmul"], tir_name
+
+
+def get_tir_w8a8_block_fp8_group_matmul(
+    N: int,
+    K: int,
+    num_experts: int,
+    block_n: int,
+    block_k: int,
+    in_dtype: Literal["float8_e4m3fn"],
+    out_dtype: Literal["float16", "bfloat16"],
+    BLOCK_SIZE_M: int,
+    BLOCK_SIZE_N: int,
+    BLOCK_SIZE_K: int,
+    GROUP_SIZE_M: int,
+    num_warps: int,
+    num_stages: int,
+    extern_mods: List[tvm.runtime.Module],  # noqa: UP006
+):
+    """Get the TIR function for the w8a8_block_fp8_group_gemm kernel."""
+    if triton is None:
+        raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.")
+
+    name_suffix = (
+        f"_N{N}_K{K}_num_experts{num_experts}_block_n{block_n}"
+        f"_block_k{block_k}_in{in_dtype}_out{out_dtype}"
+    )
+    kernel_name = f"triton_w8a8_block_fp8_group_gemm{name_suffix}"
+    tir_name = f"tir_w8a8_block_fp8_group_gemm{name_suffix}"
+    for ext_mod in extern_mods:
+        if ext_mod.implements_function(kernel_name):
+            return [None, tir_name]
+
+    triton_kernel = _get_triton_w8a8_block_fp8_group_gemm()
+    triton_kernel.__name__ = kernel_name
+
+    @I.ir_module
+    class BlockFP8GroupMatmul:
+        @T.prim_func(private=True, s_tir=True)
+        def tir_w8a8_block_fp8_group_gemm(
+            var_A: T.handle,
+            var_B: T.handle,
+            var_As: T.handle,
+            var_Bs: T.handle,
+            var_expert_ids: T.handle,
+            var_indptr: T.handle,
+            var_C: T.handle,
+        ):
+            T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1})
+            EM = T.int32()
+            A = T.match_buffer(var_A, (EM, K), dtype=in_dtype)
+            B = T.match_buffer(var_B, (num_experts, N, K), dtype=in_dtype)
+            As = T.match_buffer(var_As, (EM, (K + block_k - 1) // block_k), "float32")
+            Bs = T.match_buffer(
+                var_Bs,
+                (
+                    num_experts,
+                    (N + block_n - 1) // block_n,
+                    (K + block_k - 1) // block_k,
+                ),
+                "float32",
+            )
+            expert_ids = T.match_buffer(
+                var_expert_ids,
+                ((EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,),
+                "int32",
+            )
+            indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32")
+            C = T.match_buffer(var_C, (EM, N), dtype=out_dtype)
+
+            with T.sblock("root"):
+                T.reads(
+                    A[0:EM, 0:K],
+                    B[0:num_experts, 0:N, 0:K],
+                    As[0:EM, 0 : (K + block_k - 1) // block_k],
+                    Bs[
+                        0:num_experts,
+                        0 : (N + block_n - 1) // block_n,
+                        0 : (K + block_k - 1) // block_k,
+                    ],
+                    expert_ids[0 : (EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts],
+                    indptr[0 : num_experts + 1],
+                )
+                T.writes(C[0:EM, 0:N])
+                T.call_kernel(
+                    triton.jit(triton_kernel),
+                    ((T.ceildiv(EM, BLOCK_SIZE_M) + num_experts) * T.ceildiv(N, BLOCK_SIZE_N),),
+                    A.data,
+                    B.data,
+                    C.data,
+                    As.data,
+                    Bs.data,
+                    expert_ids.data,
+                    indptr.data,
+                    EM,
+                    N,
+                    K,
+                    num_experts,
+                    K,  # stride_am
+                    1,  # stride_ak
+                    N * K,  # stride_be
+                    1,  # stride_bk
+                    K,  # stride_bn
+                    N,  # stride_cm
+                    1,  # stride_cn
+                    (K + block_k - 1) // block_k,  # stride_asm
+                    1,  # stride_ask
+                    ((N + block_n - 1) // block_n) * ((K + block_k - 1) // block_k),  # stride_bse
+                    1,  # stride_bsk
+                    (K + block_k - 1) // block_k,  # stride_Bs_n
+                    block_n,
+                    block_k,
+                    BLOCK_SIZE_M,
+                    BLOCK_SIZE_N,
+                    BLOCK_SIZE_K,
+                    GROUP_SIZE_M,
+                    K % BLOCK_SIZE_K == 0,
+                    num_warps=num_warps,
+                    num_stages=num_stages,
+                )
+
+    new_ext_mods = BlockFP8GroupMatmul.attrs["external_mods"]
+    assert len(new_ext_mods) == 1
+    extern_mods.append(new_ext_mods[0])
+    return BlockFP8GroupMatmul["tir_w8a8_block_fp8_group_gemm"], tir_name
+
+
+def _compute_expert_id_per_block(
+    indptr: nn.Tensor,
+    num_experts: int,
+    M: nn.IntExpr,
+    BLOCK_SIZE_M: int,
+) -> nn.Tensor:
+    """Compute the expert id for each threadblock (CTA).
+    We assign an expert id to each threadblock, and the threadblock will
+    compute the gemm with regard to the specified expert.
+
+    Parameters
+    ----------
+    indptr : nn.Tensor
+        The indptr tensor of group gemm, with shape of [num_experts + 1,].
+
+    num_experts : int
+        The number of total experts.
+
+    M : nn.IntExpr
+        The number of tokens.
+
+    BLOCK_SIZE_M : int
+        The block size of the threadblock along the batch dimension.
+
+    Returns
+    -------
+    expert_ids : nn.Tensor
+        The expert id for each threadblock, with shape of
+        [(M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,].
+    """
+
+    @T.prim_func(s_tir=True)
+    def tir_compute_expert_id_per_block(
+        var_indptr: T.handle,
+        var_expert_ids: T.handle,
+        M: T.int64,
+    ):
+        T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1})
+        indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32")
+        expert_ids = T.match_buffer(
+            var_expert_ids,
+            ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,),
+            "int32",
+        )
+        with T.sblock("root"):
+            for eid in T.thread_binding(0, num_experts, thread="threadIdx.x"):
+                start_block_id: T.int32 = (indptr[eid] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid
+                num_blocks: T.int32 = (
+                    indptr[eid + 1] - indptr[eid] + BLOCK_SIZE_M - 1
+                ) // BLOCK_SIZE_M
+                start_block_id_next_expert: T.int32 = (
+                    (indptr[eid + 1] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid + 1
+                )
+                for block_id in T.serial(num_blocks):
+                    expert_ids[start_block_id + block_id] = eid
+                for block_id in T.serial(
+                    start_block_id_next_expert - (start_block_id + num_blocks)
+                ):
+                    expert_ids[start_block_id + num_blocks + block_id] = -1
+
+    assert num_experts <= 1024
+    return nn.tensor_ir_op(
+        tir_compute_expert_id_per_block,
+        "tir_compute_expert_id_per_block",
+        args=[indptr, M],
+        out=nn.Tensor.placeholder(
+            ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), dtype="int32"
+        ),
+    )
+
+
+def fp8_groupwise_scaled_gemm(
+    x: nn.Tensor,
+    x_scale: nn.Tensor,
+    weight: nn.Tensor,
+    weight_scale: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+) -> nn.Tensor:
+    """Triton block-scale fp8 gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    x_scale : nn.Tensor
+        The scale tensor, with shape of [m, k // block_size].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [n, k].
+
+    weight_scale : nn.Tensor
+        The scale tensor, with shape of [n // block_size, k // block_size].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    out_dtype : str
+        The data type of the output tensor.
+
+    Returns
+    -------
+    out : nn.Tensor
+        The output tensor, with shape of [m, n] and dtype of `out_dtype`.
+    """
+    assert x.ndim >= 2
+    assert weight.ndim == 2
+    assert x_scale.ndim == x.ndim
+    assert weight_scale.ndim == weight.ndim
+    assert x.shape[-1] == weight.shape[1]
+    assert x.shape[:-1] == x_scale.shape[:-1]
+    assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1]
+    assert (weight.shape[1] + block_size[1] - 1) // block_size[1] == weight_scale.shape[1]
+    assert (weight.shape[0] + block_size[0] - 1) // block_size[0] == weight_scale.shape[0]
+
+    if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn":
+        raise ValueError(
+            f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}"
+        )
+    if x_scale.dtype != "float32" and weight_scale.dtype != "float32":
+        raise ValueError(
+            "x_scale and weight_scale must be float32, but got "
+            f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}"
+        )
+    if out_dtype not in ["float16", "bfloat16"]:
+        raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}")
+
+    M = x.shape[0]
+    for i in range(1, x.ndim - 1):
+        M *= x.shape[i]
+    N = weight.shape[0]
+    K = x.shape[-1]
+
+    BLOCK_SIZE_M = 64
+    BLOCK_SIZE_N = block_size[0]
+    BLOCK_SIZE_K = block_size[1]
+    GROUP_SIZE_M = 32
+    num_warps = 4
+    num_stages = 3
+
+    x_shape = x.shape
+    if x.ndim > 2:
+        x = x.reshape(M, K)
+    x_scale = x_scale.reshape(M, x_scale.shape[-1])
+
+    out = nn.extern(
+        "mlc.triton.w8a8_block_fp8_matmul",
+        args=[
+            x,
+            weight,
+            x_scale,
+            weight_scale,
+            N,
+            K,
+            block_size[0],
+            block_size[1],
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            str(x.dtype),
+            str(out_dtype),
+        ],
+        out=nn.Tensor.placeholder((M, N), dtype=out_dtype),
+    )
+    return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out
+
+
+def fp8_groupwise_scaled_group_gemm(
+    x: nn.Tensor,
+    x_scale: nn.Tensor,
+    weight: nn.Tensor,
+    weight_scale: nn.Tensor,
+    indptr: nn.Tensor,
+    block_size: Tuple[int, int],  # noqa: UP006
+    out_dtype: str,
+):
+    """Triton block-scale fp8 group gemm operator.
+
+    Parameters
+    ----------
+    x : nn.Tensor
+        The input tensor, with shape of [m, k].
+
+    x_scale : nn.Tensor
+        The scale tensor, with shape of [m, k // block_size].
+
+    weight : nn.Tensor
+        The weight tensor, with shape of [num_experts, n, k].
+
+    weight_scale : nn.Tensor
+        The scale tensor, with shape of [num_experts, n // block_size, k // block_size].
+
+    indptr : nn.Tensor
+        The indptr tensor of group gemm, with shape of [num_experts + 1,].
+
+    block_size : Tuple[int, int]
+        The block size.
+
+    out_dtype : str
+        The data type of the output tensor.
+
+    Returns
+    -------
+    out : nn.Tensor
+        The output tensor, with shape of [m, n] and dtype of `out_dtype`.
+    """
+    assert x.ndim >= 2
+    assert weight.ndim == 3
+    assert x_scale.ndim == x.ndim
+    assert weight_scale.ndim == weight.ndim
+    assert x.shape[-1] == weight.shape[2]
+    assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1]
+    assert (weight.shape[2] + block_size[1] - 1) // block_size[1] == weight_scale.shape[2]
+    assert (weight.shape[1] + block_size[0] - 1) // block_size[0] == weight_scale.shape[1]
+
+    num_experts = weight.shape[0]
+    M = x.shape[0]
+    for i in range(1, x.ndim - 1):
+        M *= x.shape[i]
+    N = weight.shape[1]
+    K = x.shape[-1]
+    assert weight_scale.shape[0] == num_experts
+    assert indptr.ndim == 1
+    assert indptr.shape[0] == num_experts + 1
+
+    BLOCK_SIZE_M = 64
+    BLOCK_SIZE_N = block_size[0]
+    BLOCK_SIZE_K = block_size[1]
+    GROUP_SIZE_M = 32
+    num_warps = 4
+    num_stages = 3
+
+    x_shape = x.shape
+    if x.ndim > 2:
+        x = x.reshape(M, K)
+    x_scale = x_scale.reshape(M, x_scale.shape[-1])
+    expert_ids = _compute_expert_id_per_block(indptr, num_experts, M, BLOCK_SIZE_M)
+
+    out = nn.extern(
+        "mlc.triton.w8a8_block_fp8_group_matmul",
+        args=[
+            x,
+            weight,
+            x_scale,
+            weight_scale,
+            expert_ids,
+            indptr,
+            N,
+            K,
+            num_experts,
+            block_size[0],
+            block_size[1],
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            str(x.dtype),
+            str(out_dtype),
+        ],
+        out=nn.Tensor.placeholder((M, N), dtype=out_dtype),
+    )
+    return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out