chore: import upstream snapshot with attribution

python/tvm/relax/frontend/__init__.py

@@ -0,0 +1,21 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Frontends for constructing Relax programs, with the model importers"""
+
+from . import nn
+from .common import detach_params

python/tvm/relax/frontend/common.py

@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""Commons for Relax frontend."""
+
+import numpy as _np
+
+import tvm
+from tvm import topi
+
+
+def detach_params(mod: tvm.IRModule) -> tuple[tvm.IRModule, dict[str, list[tvm.runtime.Tensor]]]:
+    """Detach the attribute "params" in the functions of the input IRModule as
+    separate dictionary of params.
+
+    Parameters
+    ----------
+    mod : tvm.IRModule
+        The IRModule whose functions' "param" attribute is going to be detached.
+
+    Returns
+    -------
+    detached_mod : tvm.IRModule
+        The IRModule after the detachment.
+
+    params_dict : Dict[str, List[tvm.runtime.Tensor]]
+        The detached params. The dict keys corresponds to the names of the
+        functions in the input IRModule that have attribute "params".
+    """
+    detached_mod = tvm.IRModule()
+    params_dict = dict()
+    for gv, func in mod.functions_items():
+        if "params" in func.attrs:
+            params = list(func.attrs["params"])
+            if not all([isinstance(param, tvm.runtime.Tensor) for param in params]):
+                raise ValueError('The value "params" attribute is expected to be a list of Tensor.')
+            params_dict[gv.name_hint] = params
+            detached_mod[gv] = func.without_attr("params")
+        else:
+            detached_mod[gv] = func
+    return detached_mod, params_dict
+
+
+def autopad(
+    bb,
+    data,
+    strides,
+    kernel_shape,
+    dilations=(1, 1),
+    pad_type="constant",
+    deconv=False,
+    mode="SAME_UPPER",
+    pad_value=0.0,
+):
+    """
+    Perform autopadding with dynamic input shapes
+    """
+    # get attributes as constants
+    strides = _np.array(strides)
+    dilated_kernel_shape = _np.array(
+        [(kernel - 1) * dilation + 1 for kernel, dilation in zip(kernel_shape, dilations)]
+    )
+    # get input shape
+    ndim = data.ty.ndim
+    data_shape = list(data.ty.shape)
+    shape = data_shape[2:ndim]
+
+    # set up integer constants
+    zero = 0
+    one = 1
+    two = 2
+
+    # Calculate total padding
+    mod = shape % strides
+
+    left = _np.maximum(dilated_kernel_shape - strides, zero)
+    right = _np.maximum(dilated_kernel_shape - mod, zero)
+
+    total_pad = _np.where(_np.equal(mod, zero), left, right)
+    if deconv:
+        total_pad = _np.array(kernel_shape) - one - total_pad
+
+    # split total padding into before and after
+    pad_before = _np.floor_divide(total_pad, two)
+    pad_after = total_pad - pad_before
+
+    # combine
+    if "LOWER" in mode:
+        pad = _np.concatenate(
+            [_np.reshape(pad_after, [-1, 1]), _np.reshape(pad_before, [-1, 1])], axis=1
+        )
+    else:
+        pad = _np.concatenate(
+            [_np.reshape(pad_before, [-1, 1]), _np.reshape(pad_after, [-1, 1])], axis=1
+        )
+
+    # pad N and C with zeros
+    pad = _np.concatenate([_np.zeros([2, 2], dtype="int64"), pad], axis=0)
+
+    if pad_type not in ["constant", "edge", "reflect"]:
+        raise tvm.error.OpAttributeInvalid(
+            "Value " + pad_type + ' in attribute "mode" is invalid for operator Pad.'
+        )
+
+    if pad_type == "constant":
+        return bb.emit_te(topi.nn.pad, data, pad[:, 0].tolist(), pad[:, 1].tolist(), pad_value)
+    elif pad_type == "reflect":
+        return bb.emit_te(
+            topi.nn.mirror_pad, data, pad[:, 0].tolist(), pad[:, 1].tolist(), "REFLECT"
+        )
+    else:
+        # edge mode - replicate border values
+        return bb.emit_te(topi.nn.replicate_pad, data, pad[:, 0].tolist(), pad[:, 1].tolist())

python/tvm/relax/frontend/nn/__init__.py

@@ -0,0 +1,53 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""A PyTorch-like API to build IRModules."""
+
+# pylint: disable=redefined-builtin
+from . import op, spec
+from .core import (
+    Effect,
+    Module,
+    ModuleDict,
+    ModuleList,
+    Object,
+    Parameter,
+    ParameterDict,
+    ParameterList,
+    Tensor,
+)
+from .exporter import add_extern
+from .extern import ExternModule, ObjectModule, SourceModule
+from .modules import (
+    GELU,
+    Conv1D,
+    Conv2D,
+    Conv3D,
+    ConvTranspose1D,
+    Embedding,
+    GroupNorm,
+    IOEffect,
+    KVCache,
+    LayerNorm,
+    Linear,
+    ReLU,
+    RMSNorm,
+    SiLU,
+)
+from .op import *
+from .subroutine import SubroutineMixin
+from .visitor import Mutator

python/tvm/relax/frontend/nn/_tensor_op.py

@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+"""Adding member operators to nn.Tensor."""
+
+from tvm import tirx
+
+
+def _op():
+    from tvm.relax.frontend.nn import op  # pylint: disable=import-outside-toplevel
+
+    return op
+
+
+def _convert_scalar(scalar, ref) -> "Tensor":
+    from .core import Tensor  # pylint: disable=import-outside-toplevel
+
+    if isinstance(scalar, Tensor):
+        return scalar
+    if isinstance(scalar, tirx.FloatImm | tirx.IntImm):
+        return Tensor.from_scalar(scalar.value, dtype=ref.dtype)
+    if isinstance(scalar, int | float):
+        return Tensor.from_scalar(scalar, dtype=ref.dtype)
+    return scalar
+
+
+class _TensorOp:
+    def __add__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().add(self, other)
+
+    def __radd__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().add(self, other)
+
+    def __sub__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().subtract(self, other)
+
+    def __rsub__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().subtract(other, self)
+
+    def __mul__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().multiply(self, other)
+
+    def __rmul__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().multiply(self, other)
+
+    def __truediv__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().divide(self, other)
+
+    def __lt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less(self, other)
+
+    def __le__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less_equal(self, other)
+
+    def __gt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater(self, other)
+
+    def __ge__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater_equal(self, other)
+
+    def astype(self, dtype):
+        return _op().astype(self, dtype)
+
+    def maximum(self, other):
+        other = _convert_scalar(other, self)
+        return _op().maximum(self, other)
+
+    def minimum(self, other):
+        other = _convert_scalar(other, self)
+        return _op().minimum(self, other)
+
+    def reshape(self, *shape):
+        return _op().reshape(self, shape)
+
+    def permute_dims(self, *axes):
+        return _op().permute_dims(self, axes)
+
+    def repeat(self, repeats: int, axis: int | None = None):
+        return _op().repeat(self, repeats, axis)

python/tvm/relax/frontend/nn/core.py

@@ -0,0 +1,879 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The core infra for nn.Module, which includes the following pieces:
+- Tensor, a wrapper on top of relax.Expr whose ty is a TensorType,
+  providing more convenient access shape and dtype information.
+  Tensor is always symbolic and not bound to any concrete values.
+- Parameter, a special tensor which could be bound or not bound to concrete values.
+- Module, a container of nn.Parameters and sub nn.Modules.
+- Effect, a non-user-facing class that encloses potential side effects, for example, IO,
+  impure external function callings, inplace mutation, etc.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable, Iterator, Sequence
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Union,
+)
+
+import numpy as np  # type: ignore
+
+import tvm.runtime
+from tvm import tirx
+from tvm.ir import IRModule
+from tvm.ir.transform import Pass
+from tvm.runtime import Device
+from tvm.runtime import device as as_device
+from tvm.runtime.vm import VirtualMachine
+from tvm.target import Target
+
+from .... import relax as rx
+from ...block_builder import BlockBuilder
+from ...type import (
+    AnyType,
+    ShapeType,
+    TensorType,
+    TupleType,
+)
+from ._tensor_op import _TensorOp
+from .subroutine import SubroutineMixin
+
+if TYPE_CHECKING:
+    import torch  # type: ignore
+
+    from . import spec as _spec
+    from .extern import ExternModule
+
+
+_DEFAULT_DTYPE = "float32"
+
+
+def get_default_dtype() -> str:
+    """Get the default parameter dtype if not specified. By default it is float32.
+
+    Returns
+    -------
+    dtype : str
+        The default dtype
+    """
+    return _DEFAULT_DTYPE
+
+
+def set_default_dtype(dtype: str) -> None:
+    """Set the default parameter dtype.
+
+    Parameters
+    ----------
+    dtype : str
+        The default dtype to be set
+    """
+    global _DEFAULT_DTYPE  # pylint: disable=global-statement
+    _DEFAULT_DTYPE = dtype
+
+
+class Tensor(_TensorOp):
+    """A wrapper on top of relax.Expr whose ty is a TensorType, providing more
+    convenient access shape and dtype information. Tensor is always symbolc and not bound to any
+    concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when
+    operators are applied on tensors, the shape and dtype information is already available.
+    """
+
+    _expr: rx.Expr
+
+    def __init__(self, *, _expr: rx.Expr) -> None:
+        """Private constructor. Tensor is never supposed to be constructed directly by users."""
+
+        def _check_tensor(expr: rx.Expr) -> None:
+            assert expr.ty is not None
+            assert isinstance(expr.ty, TensorType)
+            assert expr.ty.ndim != -1
+            assert expr.ty.shape is not None
+            assert expr.ty.shape.ty is not None
+            assert isinstance(expr.ty.shape.ty, ShapeType)
+            assert expr.ty.shape.ty.values is not None
+
+        _check_tensor(_expr)
+        self._expr = _expr
+
+    @staticmethod
+    def from_const(data) -> "Tensor":
+        """Construct a tensor from numpy constants."""
+        return Tensor(_expr=rx.const(data))
+
+    @staticmethod
+    def from_scalar(data: int | float, dtype: str) -> "Tensor":
+        """Construct a tensor from a scalar with dtype specified."""
+        return Tensor(_expr=rx.const(data, dtype=dtype))
+
+    @staticmethod
+    def from_ty(ty: rx.TensorType, name: str = "tensor") -> "Tensor":
+        """Construct a nn.Tensor from a Relax TensorType.
+
+        TensorType is the Relax type-level description of a tensor, carrying its shape
+        and dtype without holding actual data. This factory creates an unbound placeholder
+        ``nn.Tensor`` that can be used as a symbolic input when tracing an ``nn.Module``.
+
+        Parameters
+        ----------
+        ty : rx.TensorType
+            The type describing the tensor's shape and dtype.
+
+        name : str
+            Name hint for the underlying Relax variable.
+
+        Returns
+        -------
+        tensor : Tensor
+            A symbolic ``nn.Tensor`` backed by a ``relax.Var`` with the given type.
+        """
+        return Tensor(
+            _expr=rx.Var(
+                name_hint=name,
+                ty=ty,
+            )
+        )
+
+    @staticmethod
+    def placeholder(
+        shape: Sequence[int | str | tirx.Expr],
+        dtype: str,
+        name: str = "tensor",
+    ) -> "Tensor":
+        """Create a placeholder tensor with given shape and dtype. A placeholder tensor should
+        never be created directly by users in usual cases, and the only exception is to indicate
+        the shape/dtype of return values of an external function.
+
+        If shape is a string `name`, we create a symbolic shape `tvm.tirx.Var(name, "int64")`.
+        """
+        new_shape = []
+        for expr in shape:
+            if isinstance(expr, int | tirx.IntImm):
+                expr = int(expr)
+                assert expr >= 0
+                new_shape.append(expr)
+                continue
+            if isinstance(expr, str):
+                expr = tirx.Var(expr, "int64")
+                new_shape.append(expr)
+                continue
+            if not tvm.ir.is_prim_expr(expr):
+                raise TypeError(f"Invalid shape: {shape}")
+            assert expr.ty == tvm.ir.PrimType("int64")
+            new_shape.append(expr)
+        return Tensor(
+            _expr=rx.Var(
+                name_hint=name,
+                ty=TensorType(
+                    shape=new_shape,  # type: ignore[arg-type]
+                    dtype=dtype,
+                ),
+            )
+        )
+
+    @property
+    def shape(self) -> list[int | tirx.Expr]:
+        """Returns the shape of the tensor as a list of integers.
+
+        An integer can be a python int or tvm.tirx.Expr, depending on whether the shape is
+        fully static, for example, [1, 2, tvm.tirx.Var("n")] is a valid shape where the last
+        dimension is dynamic while the first two dimensions are always static constants.
+
+        Returns
+        -------
+        shape : List[Union[int, tirx.Expr]]
+            The shape of the tensor
+        """
+
+        def _simplify(expr: tirx.Expr):
+            return expr.value if isinstance(expr, tirx.IntImm) else expr
+
+        shape_ty: ShapeType = self._expr.ty.shape.ty
+        return [_simplify(x) for x in shape_ty.values]
+
+    @property
+    def ndim(self) -> int:
+        """Returns the number of dimensions of the tensor.
+
+        Returns
+        -------
+        ndim : int
+            The number of dimensions of the tensor
+        """
+        return self._expr.ty.ndim
+
+    @property
+    def dtype(self) -> str:
+        """Returns the data type of the tensor.
+
+        Returns
+        -------
+        dtype : str
+            The data type of the tensor
+        """
+        return self._expr.ty.dtype
+
+    def __repr__(self) -> str:
+        return f'Tensor({self.shape}, "{self.dtype}")'
+
+
+class Parameter(Tensor):
+    """A parameter represents the weight of a neural network layer. It is a special tensor which
+    could be bound or not bound to concrete values. If a parameter is bound to a concrete value,
+    it is called a bound parameter, otherwise it is called an unbound parameter.
+    """
+
+    _data: Tensor | None
+    attrs: dict[str, Any]
+
+    def __init__(
+        self,
+        shape: Sequence[int | str | tirx.Expr],
+        dtype: str | None = None,
+    ) -> None:
+        """Create a parameter with given shape and dtype. The parameter is not bound to any
+        concrete values.
+
+        Parameters
+        ----------
+        shape : Sequence[Union[int, str, tirx.Expr]]
+            The shape of the parameter. If it is a string `name`, we create a symbolic shape
+            `tvm.tirx.Var(name, "int64")`.
+        dtype : Optional[str]
+            The data type of the parameter. If not specified, the default dtype will be used.
+        """
+        if dtype is None:
+            dtype = get_default_dtype()
+        super().__init__(_expr=Tensor.placeholder(shape, dtype=dtype, name="param")._expr)
+        self._data = None
+        self.attrs = OrderedDict()
+
+    @property
+    def data(self) -> Tensor | None:
+        """Returns the concrete value of the parameter if it is bound to a concrete value,
+        otherwise returns None. The returned value is a tvm.runtime.Tensor."""
+        return self._data
+
+    @data.setter
+    def data(self, data: Union[None, tvm.runtime.Tensor, np.ndarray, "torch.Tensor"]) -> None:
+        """Set the concrete value of the parameter. The data should be one of the following:
+        - None: unbind the parameter to concrete values
+        - tvm.runtime.Tensor
+        - numpy.ndarray
+        - torch.Tensor and any other DLPack-compliant tensors
+        """
+        if data is None:
+            self._data = data
+            return
+        # Try to do zero-copy if possible
+        if isinstance(data, tvm.runtime.Tensor):
+            pass
+        elif isinstance(data, np.ndarray):
+            data = tvm.runtime.tensor(data)
+        elif hasattr(data, "__dlpack__"):
+            data = _from_dlpack(data)
+        else:
+            raise TypeError(f"Unsupported data type: {type(data)}")
+        if data.shape != tuple(self.shape):
+            raise ValueError(f"Shape mismatch: expected {tuple(self.shape)}, got {data.shape}")
+        if data.dtype != self.dtype:
+            raise ValueError(f"Dtype mismatch: expected {self.dtype}, got {data.dtype}")
+        self._data = data
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Change the dtype of the parameter if it is not bound to any concrete data"""
+        if dtype is not None:
+            if self._data is not None:
+                raise ValueError(
+                    "Changing the dtype of a Parameter that has been bound to concrete "
+                    "data is not recommended. It might lead to potential precision loss "
+                    "or other unexpected behaviors"
+                )
+            self._expr = Tensor.placeholder(  # pylint: disable=protected-access
+                self.shape, dtype=dtype, name="param"
+            )._expr
+
+
+class Object:
+    """A wrapper on top of relax.Expr whose ty is the base
+    AnyType, rather than a more specific subtype. Object effectively
+    represents non-tensor frontend components such as KV caches.
+    """
+
+    _expr: rx.Var
+
+    def __init__(self, *, _expr: rx.Expr, _name: str) -> None:
+        """Private constructor. Object is never supposed to be constructed directly by users."""
+        if not isinstance(_expr, rx.Var):
+            _expr = BlockBuilder.current().emit(_expr, _name)
+        self._expr = _expr
+        assert isinstance(self._expr.ty, AnyType)
+
+
+class Effect:
+    """Effect is a special non-user facing type that is used to represent operations with side
+    effects, for example, print. It is used to represent the output of a computation.
+    """
+
+    def emit_init(self, name_hint: str, builder: BlockBuilder) -> list[rx.DataflowVar]:
+        """Emit the initialization of the effect. This method is called by the compiler to
+        initialize the effect."""
+        raise NotImplementedError
+
+    def create(self, name_hint: str) -> list[rx.Var]:
+        """Create the implicit inputs to a relax.Function that represents the side effect"""
+        raise NotImplementedError
+
+    def set_state(self, state_vars: list[rx.Var]) -> None:
+        """Set the variables that represents the effect"""
+        raise NotImplementedError
+
+    def finalize(self) -> list[rx.Var]:
+        """finalize the effect as the implicit return value of a relax.Function"""
+        raise NotImplementedError
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Convert the effect to specific dtype. Usually it is no-op for most of the effects"""
+
+
+class Module(SubroutineMixin):
+    """Base class for neural network components. Subclass it to build your models.
+    Modules can nest within each other in a tree structure using regular attribute assignment."""
+
+    def named_parameters(self, prefix: str = "") -> Iterator[tuple[str, Parameter]]:
+        """This method provides an iterator over module parameters,
+        yielding both the parameter name and its corresponding value.
+
+        Parameters
+        ----------
+        prefix : str
+            Prefix to prepend to all parameter names.
+
+        Yields
+        ------
+        (str, Parameter) - Tuple containing the name and parameter
+        """
+        yield from _attribute_finder(
+            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
+        )
+
+    def parameters(self) -> Iterator[Parameter]:
+        """This method provides an iterator over module parameters,
+        yielding only the Parameter value.
+
+        Yields
+        ------
+        Parameter - The module's parameter
+        """
+        for _, param in self.named_parameters():
+            yield param
+
+    def state_dict(
+        self, *, prefix: str = "", destination: dict[str, Parameter] | None = None
+    ) -> dict[str, Parameter]:
+        """Returns a dictionary containing references to the whole state of the module.
+
+        Parameters
+        ----------
+        prefix : str
+            Prefix to prepend to all parameter names.
+        destination : Optional[Dict[str, Parameter]]
+            Dictionary to which state will be saved. If None, a new dictionary is created.
+
+        Returns
+        -------
+        dict : Dict[str, Parameter]
+            a dictionary containing a whole state of the module
+        """
+        if destination is None:
+            destination = OrderedDict()
+        for name, param in _attribute_finder(
+            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
+        ):
+            destination[name] = param
+        return destination
+
+    def load_state_dict(
+        self, state_dict: dict[str, Parameter], strict: bool = True
+    ) -> tuple[list[str], list[str]]:
+        """This function copies parameters and buffers from the state_dict into the current module
+        and its descendants. If `strict` is set to True, the keys in the `state_dict` must exactly
+        match the keys returned by the `state_dict()` function of this module.
+
+        Parameters
+        ----------
+        state_dict : Dict[str, Parameter]
+            A dictionary containing a whole state of the module
+        strict : bool = True
+            Whether to strictly enforce that the keys in `state_dict` match the keys returned by
+            this module's `state_dict()` function.
+
+        Returns
+        -------
+        (missing_keys, unexpected_keys) : Tuple[List[str], List[str]]
+            A tuple of two lists: the missing keys and the unexpected keys.
+        """
+        self_state_dict = self.state_dict()
+        missing_keys: list[str] = []
+        unexpected_keys: list[str] = []
+        for key, value in state_dict.items():
+            if key not in self_state_dict:
+                unexpected_keys.append(key)
+                continue
+            if value.data is None:
+                raise ValueError(f"Parameter {key} is not set to any concrete tensor")
+            self_state_dict.pop(key).data = value.data
+        missing_keys = list(self_state_dict.keys())
+        if strict and (missing_keys or unexpected_keys):
+            raise KeyError(f"Missing keys: {missing_keys}, Unexpected keys: {unexpected_keys}")
+        return missing_keys, unexpected_keys
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the module with the given inputs and returns the output."""
+        if not hasattr(self, "forward"):
+            raise NotImplementedError(f"Module {type(self)} does not have a `forward` method")
+        return self.forward(*args, **kwargs)  # pylint: disable=no-member
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Convert the module to specific dtype recursively"""
+        for _, item in self.__dict__.items():
+            if hasattr(item, "to") and callable(item.to):
+                item.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype  # pylint: disable=attribute-defined-outside-init
+
+    def export_tvm(
+        self,
+        spec: "_spec.ModuleSpecType",
+        debug: bool = False,
+        allow_extern: bool = False,
+    ) -> (
+        tuple[IRModule, list[tuple[str, Parameter]]]
+        | tuple[IRModule, list[tuple[str, Parameter]], list["ExternModule"]]
+    ):
+        """Export the module to TVM IRModule and parameters
+
+        Parameters
+        ----------
+        spec : _spec.ModuleSpecType
+            A dictionary mapping each input name to a specification
+            that defines the inputs shape and dtype.
+        debug : bool
+            If set to True, then the exported module will support
+            effects. This enables things like printing in the graph.
+
+        Returns
+        -------
+        irmodule : tvm.ir.IRModule
+            The converted tvm IR representation of the model.
+        params : List[Tuple[str, Parameter]]
+            A list of Parameters corresponding to the weights of the model.
+        ext_mods : List[nn.ExternModule]
+            A list of ExternModules that are used in the model.
+        """
+        # pylint: disable=import-outside-toplevel
+        from . import spec as _spec
+        from .exporter import Exporter
+
+        # pylint: enable=import-outside-toplevel
+        spec = _spec.ModuleSpec.from_raw(spec, self)
+        mod, params, ext_mods = Exporter(debug=debug).build(spec)
+        if allow_extern:
+            return mod, params, ext_mods
+        if ext_mods:
+            raise ValueError(
+                "`ExternModule`(s) exist when they are not allowed. "
+                "Turn on flag `allow_extern` to allow."
+            )
+        return mod, params
+
+    def jit(  # pylint: disable=too-many-arguments
+        self,
+        spec: "_spec.ModuleSpec",
+        device: str | Device = "cpu",
+        pipeline: None | str | Pass = "default_build",
+        out_format: str = "torch",
+        debug: bool = False,
+    ) -> Any:
+        """Just-in-time compile an ``nn.Module`` into a callable executable.
+
+        The method exports the module to a Relax IRModule, applies the given compilation
+        pipeline, builds a Relax VM executable, and wraps the result so it can be called
+        directly (e.g. with PyTorch tensors when ``out_format="torch"``).
+
+        Parameters
+        ----------
+        spec : _spec.ModuleSpec
+            A specification mapping each module input to its shape and dtype.
+
+        device : Union[str, Device]
+            The device to compile and run on (e.g. ``"cpu"``, ``"cuda"``).
+
+        pipeline : Union[None, str, Pass]
+            The Relax compilation pipeline to apply. ``"default_build"`` uses the standard
+            optimization pipeline; ``None`` skips pipeline passes.
+
+        out_format : str
+            Output wrapper format. ``"torch"`` returns a ``TorchModule`` whose ``forward``
+            accepts and returns PyTorch tensors.
+
+        debug : bool
+            If ``True``, enable effect-based debugging (e.g. printing) in the compiled graph.
+
+        Returns
+        -------
+        module : Any
+            A callable wrapper (type depends on *out_format*) around the compiled VM.
+        """
+
+        def _compile(spec, device, pipeline, debug):
+            # pylint: disable=import-outside-toplevel
+            from ...transform import AttachExternModules
+            from ...vm_build import build as relax_build
+            from . import spec as _spec
+            from .exporter import Exporter
+
+            # pylint: enable=import-outside-toplevel
+
+            spec = _spec.ModuleSpec.from_raw(spec, self)
+            mod, params, ext_mods = Exporter(debug=debug).build(spec)
+            mod = AttachExternModules(ext_mods)(mod)  # pylint: disable=not-callable
+            vm = VirtualMachine(  # pylint: disable=invalid-name
+                relax_build(
+                    mod,
+                    target=Target.from_device(device),
+                    relax_pipeline=pipeline,
+                ),
+                device,
+            )
+            params = _param_to_tensor(params, device)
+            return spec, vm, params
+
+        device = as_device(device)
+        spec, vm, params = _compile(spec, device, pipeline, debug)  # pylint: disable=invalid-name
+
+        if out_format == "torch":
+            from . import torch  # pylint: disable=import-outside-toplevel
+
+            return torch.TorchModule(spec=spec, params=params, vm=vm)
+
+        raise ValueError(f"Unknown out_format: {out_format}")
+
+
+class ModuleDict(Module):
+    """Holds submodules in a dict."""
+
+    def __init__(self, modules: OrderedDict[str, Module] | None = None):
+        if modules is None:
+            self.modules = OrderedDict()
+        else:
+            self.modules = OrderedDict(modules)
+
+    def __iter__(self):
+        return iter(self.modules.values())
+
+    def __getitem__(self, key: str) -> Module:
+        return self.modules[key]
+
+    def __setitem__(self, key: str, module: Module) -> None:
+        self.modules[key] = module
+
+    def __len__(self) -> int:
+        return len(self.modules)
+
+    def keys(self) -> Iterator[str]:
+        return self.modules.keys()
+
+    def values(self) -> Iterator[Module]:
+        return self.modules.values()
+
+    def items(self) -> Iterator[tuple[str, Module]]:
+        return self.modules.items()
+
+    def get(self, key: str, default: Module | None = None) -> Module | None:
+        return self.modules.get(key, default)
+
+    def update(self, modules: dict[str, Module]) -> None:
+        self.modules.update(modules)
+
+    def clear(self) -> None:
+        self.modules.clear()
+
+    def pop(self, key: str) -> Module:
+        return self.modules.pop(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.modules
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for module in self.modules.values():
+            module.to(dtype=dtype)
+
+
+class ParameterDict(Module):
+    """Holds parameters in a dict."""
+
+    def __init__(
+        self,
+        params: OrderedDict[str, Parameter] | dict[str, Parameter] | None = None,
+    ):
+        self.params: OrderedDict[str, Parameter] = OrderedDict()
+        if params is not None:
+            self.update(params)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.params)
+
+    def __getitem__(self, key: str) -> Parameter:
+        return self.params[key]
+
+    def __setitem__(self, key: str, param: Parameter) -> None:
+        if not isinstance(key, str):
+            raise TypeError(f"ParameterDict keys must be strings, but got {type(key).__name__}")
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterDict values must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params[key] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def keys(self) -> Iterator[str]:
+        return self.params.keys()
+
+    def values(self) -> Iterator[Parameter]:
+        return self.params.values()
+
+    def items(self) -> Iterator[tuple[str, Parameter]]:
+        return self.params.items()
+
+    def get(self, key: str, default: Parameter | None = None) -> Parameter | None:
+        return self.params.get(key, default)
+
+    def update(self, params: dict[str, Parameter]) -> None:
+        for key, param in params.items():
+            self[key] = param
+
+    def clear(self) -> None:
+        self.params.clear()
+
+    def pop(self, key: str) -> Parameter:
+        return self.params.pop(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.params
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params.values():
+            param.to(dtype=dtype)
+
+
+class ModuleList(Module):
+    """Holds submodules in a list."""
+
+    def __init__(self, modules: list[Module]):
+        self.modules = modules
+
+    def __iter__(self):
+        return iter(self.modules)
+
+    def __getitem__(self, idx: int) -> Module:
+        return self.modules[idx]
+
+    def __setitem__(self, idx: int, module: Module) -> None:
+        self.modules[idx] = module
+
+    def __len__(self):
+        return len(self.modules)
+
+    def append(self, module: Module):
+        """Add a module to the end of the ModuleList"""
+        self.modules.append(module)
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for module in self.modules:
+            module.to(dtype=dtype)
+
+    def forward(self, x):  # pylint: disable=invalid-name
+        """Feed-forward pass of the module"""
+        for module in self.modules:
+            x = module(x)
+        return x
+
+
+class ParameterList(Module):
+    """Holds parameters in a list."""
+
+    def __init__(self, params: list[Parameter] | None = None):
+        self.params: list[Parameter] = []
+        if params is not None:
+            self.extend(params)
+
+    def __iter__(self) -> Iterator[Parameter]:
+        return iter(self.params)
+
+    def __getitem__(self, idx: int) -> Parameter:
+        return self.params[idx]
+
+    def __setitem__(self, idx: int, param: Parameter) -> None:
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params[idx] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def append(self, param: Parameter) -> None:
+        """Add a parameter to the end of the ParameterList"""
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params.append(param)
+
+    def extend(self, params: list[Parameter]) -> None:
+        """Add parameters to the end of the ParameterList"""
+        for param in params:
+            self.append(param)
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params:
+            param.to(dtype=dtype)
+
+
+def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
+    """Wrap the given relax.Expr, emit it using the current BlockBuilder,
+    and automatically handle nested cases if the expr represents a Tuple.
+
+    Parameters
+    ----------
+    expr : relax.Expr
+        The Expr to be wrapped.
+
+    name : str
+        Name hint.
+
+    Returns
+    -------
+    result : Union[Tensor, Tuple[Tensor]]
+        The computed result.
+    """
+    if not isinstance(expr, rx.DataflowVar):
+        expr = BlockBuilder.current().emit(expr, name)
+    if isinstance(expr.ty, TensorType):
+        return Tensor(_expr=expr)
+    if isinstance(expr.ty, TupleType):
+        return tuple(
+            wrap_nested(  # type: ignore
+                rx.TupleGetItem(expr, i),
+                name=f"{name}.{i}",
+            )
+            for i in range(len(expr.ty.fields))
+        )
+    raise TypeError(f"Unsupported return type: {expr.ty}")
+
+
+def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):
+    """Find attributes that satisfy the condition recursively"""
+    if isinstance(root, ParameterList):
+        for i, param in enumerate(root):
+            if condition_yield(param):
+                yield prefix + f"{i}", param
+        return
+    elif isinstance(root, ParameterDict):
+        for name, param in root.items():
+            if condition_yield(param):
+                yield prefix + name, param
+        return
+    elif isinstance(root, ModuleList):
+        for i, subitem in enumerate(root):
+            yield from _attribute_finder(subitem, prefix + f"{i}.", condition_yield)
+        return
+    elif isinstance(root, ModuleDict):
+        for name, subitem in root.items():
+            yield from _attribute_finder(subitem, prefix + f"{name}.", condition_yield)
+        return
+    for name, item in root.__dict__.items():
+        if condition_yield(item):
+            yield prefix + name, item
+        elif isinstance(item, ParameterList):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ParameterDict):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ModuleList):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ModuleDict):
+            for sub_name, sub_item in item.items():
+                yield from _attribute_finder(
+                    sub_item,
+                    prefix + name + f".{sub_name}.",
+                    condition_yield,
+                )
+        elif isinstance(item, Module):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+
+
+def _from_dlpack(tensor) -> tvm.runtime.Tensor:
+    try:
+        return tvm.runtime.from_dlpack(tensor)
+    except RuntimeError:
+        pass
+    # special logic for PyTorch
+    device_type = tensor.device.type
+    device_id = tensor.device.index or 0
+    return tvm.runtime.tensor(
+        tensor.numpy(),
+        device=Device(
+            Device._DEVICE_NAME_TO_TYPE[device_type],
+            device_id,
+        ),
+    )
+
+
+def _param_to_tensor(
+    params: list[tuple[str, Parameter]], device: Device
+) -> list[tvm.runtime.Tensor]:
+    results = []
+    missing = []
+    for name, param in params:
+        if param.data is None:
+            missing.append(name)
+        else:
+            results.append(param.data.copyto(target=device))
+    if missing:
+        raise ValueError(f"Parameters are not set to any concrete values: {', '.join(missing)}")
+    return results

python/tvm/relax/frontend/nn/exporter.py

@@ -0,0 +1,334 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Export `nn.Module` to TVM's IRModule."""
+
+import functools
+import operator
+import threading
+import typing
+
+from tvm import tirx
+from tvm.ir import IRModule
+
+from .... import relax as rx
+from ...block_builder import BlockBuilder
+from ...type import AnyType, ShapeType, TupleType
+from . import core, extern
+from . import spec as _spec
+from .modules import IOEffect
+
+
+def add_extern(mod: extern.ExternModule) -> None:
+    """Add an external module to the exporter."""
+    try:
+        exporter = Exporter.current()
+    except Exception as exception:
+        raise RuntimeError(
+            "`nn.add_extern` should only be invoked when exporting a module."
+        ) from exception
+    exporter.add_external_module(mod)
+
+
+class Exporter:
+    """Builder of ModuleSpec, which exports an nn.Module to TVM IRModule."""
+
+    _tls = threading.local()
+
+    builder: BlockBuilder
+    io_effect: core.Effect
+    extern_mods: list[extern.ExternModule]
+
+    def __init__(self, debug: bool) -> None:
+        self.builder = BlockBuilder()
+        self.io_effect = IOEffect() if debug else None
+        self.extern_mods = []
+
+    @staticmethod
+    def current() -> "Exporter":
+        """Get the current Exporter under the with scope."""
+        assert hasattr(Exporter._tls, "current")
+        return Exporter._tls.current
+
+    def __enter__(self) -> "Exporter":
+        assert not hasattr(Exporter._tls, "current")
+        Exporter._tls.current = self
+        return self
+
+    def __exit__(self, exc_type, exc, traceback) -> None:
+        assert hasattr(Exporter._tls, "current")
+        delattr(Exporter._tls, "current")
+
+    def add_external_module(self, mod: extern.ExternModule) -> None:
+        """Add an external module to the exporter."""
+        # pylint: disable=protected-access
+        all_symbols: list[str] = []
+        for extern_mod in self.extern_mods:
+            all_symbols.extend(extern_mod._symbols.keys())
+        duplicated_symbols = list(set(mod._symbols.keys()) & set(all_symbols))
+        # pylint: enable=protected-access
+        if duplicated_symbols:
+            raise ValueError(f"Duplicate symbols: {duplicated_symbols}")
+        self.extern_mods.append(mod)
+
+    def build(  # pylint: disable=too-many-locals
+        self,
+        spec: _spec.ModuleSpec,
+    ) -> tuple[
+        IRModule,
+        list[tuple[str, core.Parameter]],
+        list[extern.ExternModule],
+    ]:
+        """Build the ModuleSpec to TVM IRModule. Returns the IRModule and the parameters."""
+
+        # pylint: disable=protected-access
+        def _params() -> list[tuple[str, core.Parameter]]:
+            params = []
+            for name, param in core._attribute_finder(
+                spec.module, prefix="", condition_yield=lambda x: isinstance(x, core.Parameter)
+            ):
+                params.append((name, param))
+            return params
+
+        def _effects() -> list[tuple[str, core.Effect]]:
+            result = []
+            if self.io_effect is not None:
+                result.append(("", self.io_effect))
+            for name, effect in core._attribute_finder(
+                spec.module, "", condition_yield=lambda x: isinstance(x, core.Effect)
+            ):
+                result.append((name, effect))
+            return result
+
+        # pylint: enable=protected-access
+        params = None
+        effects = _effects()
+        ext_mods = self.extern_mods
+        with self:
+            if effects:
+                with self.builder.function("_initialize_effect"):
+                    with self.builder.dataflow():
+                        outputs = _emit_effect_init(self.builder, effects)
+                    self.builder.emit_func_output(outputs, params=[])
+            for method_name, method_spec in zip(spec.method_names, spec.method_specs):
+                params = _params()  # Re-initialize so symbolic shapes not shared across methods
+                len_args = len(method_spec.arg_specs)
+                len_effects = {
+                    "packed": 1,
+                    "none": 0,
+                    "plain": len(effects),
+                }[method_spec.effect_mode]
+                with self.builder.function(
+                    method_name,
+                    attrs={"num_input": len_args + len_effects},  # type: ignore
+                ):
+                    with self.builder.dataflow():
+                        outputs, inputs = _emit_method(self.builder, method_spec, params, effects)
+                    self.builder.emit_func_output(outputs, inputs)
+        mod = self.builder.finalize()
+        rx.analysis.well_formed(mod)
+
+        return mod, params, ext_mods
+
+
+def _emit_effect_init(
+    builder: BlockBuilder,
+    effects: list[tuple[str, core.Effect]],
+):
+    outputs = []
+    for prefix, effect in effects:
+        inits = effect.emit_init(prefix, builder)
+        assert isinstance(inits, list)
+        outputs.extend(inits)
+    outputs = builder.emit_output(builder.emit(rx.Tuple(outputs)))
+    return outputs
+
+
+def _emit_method(  # pylint: disable=too-many-locals,too-many-branches,too-many-statements
+    builder: BlockBuilder,
+    spec: _spec.MethodSpec,
+    params: list[tuple[str, core.Parameter]],
+    effects: list[tuple[str, core.Effect]] | None,
+):
+    # pylint: disable=protected-access
+    # symbolic shape's name mapping to its tirx.Var for reuse
+    str2var_params: dict[str, tirx.Var] = {}
+
+    def _unwrap_ret(expr: typing.Any) -> typing.Any:
+        if isinstance(expr, core.Tensor | core.Object):
+            return expr._expr
+        if isinstance(expr, tuple):
+            return rx.Tuple([_unwrap_ret(x) for x in expr])
+        if isinstance(expr, list):
+            return rx.Tuple([_unwrap_ret(x) for x in expr])
+        raise TypeError(f"Unsupported return type: {type(expr)}")
+
+    def _convert_input(arg):
+        if isinstance(arg, tirx.Var):
+            return rx.Var(arg.name, ty=ShapeType(values=[arg]))
+        if isinstance(arg, core.Tensor | core.Object):
+            return arg._expr  # pylint: disable=protected-access
+        if isinstance(arg, _spec.Tuple):
+            return rx.Var(
+                arg.name,
+                ty=TupleType([_convert_input(arg_i).ty for arg_i in arg.elements]),
+            )
+        raise TypeError(f"Unsupported input type: {type(arg)}")
+
+    def _params(mode: str) -> list[rx.Var]:
+        inputs: list[rx.Var] = []
+
+        def _get_var(shape_var: tirx.Var) -> tirx.Var:
+            name = shape_var.name
+            if name in str2var_params:
+                return str2var_params[name]
+            var = tirx.Var(name, "int64")
+            str2var_params[name] = var
+            return var
+
+        for name, param in params:
+            # Make sure the a symbolic shape is not re-registered (same as _method_spec_to_inputs)
+            # e.g. we do not see `vocab_size` for `lm_head` and `vocab_size_1` for `embed_tokens`
+            new_shape = [_get_var(x) if isinstance(x, tirx.Var) else x for x in param.shape]
+            var = core.Tensor.placeholder(new_shape, param.dtype, name)._expr
+            inputs.append(var)
+            param._expr = var
+        if mode == "none":
+            return []
+        if mode == "plain":
+            return inputs
+        if mode == "packed":
+            input_var = rx.Var(
+                "packed_params",
+                TupleType(fields=[x.ty for x in inputs]),
+            )
+            for i, (name, param) in enumerate(params):
+                param._expr = builder.emit(rx.TupleGetItem(input_var, i), name_hint=name)
+            return [input_var]
+        raise ValueError(f"Invalid param_mode: {mode}")
+
+    def _effects(mode: str) -> list[rx.Var]:
+        unflat_inputs: list[list[rx.Var]] = []
+        for name, effect in effects:
+            effect_input = effect.create(name)
+            effect.set_state(effect_input)
+            unflat_inputs.append(effect_input)
+        inputs: list[rx.Var] = functools.reduce(operator.iadd, unflat_inputs, [])
+        if mode == "none":
+            return []
+        if mode == "plain":
+            return inputs
+        if mode == "packed":
+            input_var = rx.Var(
+                "packed_effects",
+                TupleType(fields=[x.ty for x in inputs]),
+            )
+            i = 0
+            for effect_input, (_, effect) in zip(unflat_inputs, effects):
+                updated_effect_input = []
+                for effect_input_i in effect_input:
+                    updated_effect_input.append(
+                        builder.emit(
+                            rx.TupleGetItem(input_var, i),
+                            name_hint=effect_input_i.name_hint,
+                        )
+                    )
+                    i += 1
+                effect.set_state(updated_effect_input)
+            return [input_var]
+
+        raise ValueError(f"Invalid effect_mode: {mode}")
+
+    # pylint: enable=protected-access
+
+    def _detuple(arg, var: rx.Var, builder: BlockBuilder):
+        if isinstance(arg, _spec.Tuple):
+            ret = []
+            for i, elem in enumerate(arg.elements):
+                field = builder.emit(rx.TupleGetItem(var, i), name_hint=f"{arg.name}_{i}")
+                ret.append(_detuple(elem, field, builder))
+            return type(arg.elements)(ret)
+        if isinstance(arg, core.Tensor):
+            return core.Tensor(_expr=var)
+        if isinstance(arg, tirx.Var):
+            return arg
+        raise TypeError(f"Unsupported input type: {type(arg)}")
+
+    # TODO(@junrushao): Warn if params/effects are used when their mode is "none"
+    explicit_inputs = _method_spec_to_inputs(spec)
+    inputs = [_convert_input(x) for x in explicit_inputs]
+    inputs = inputs + _effects(spec.effect_mode)
+    inputs = inputs + _params(spec.param_mode)
+
+    for arg_idx, (arg, var) in enumerate(zip(explicit_inputs, inputs)):
+        if isinstance(arg, _spec.Tuple):
+            explicit_inputs[arg_idx] = _detuple(arg, var, builder)
+
+    outputs = spec.method(*explicit_inputs)
+    effect_outputs = []
+    for _, effect in effects:
+        effect_outputs.extend(effect.finalize())
+    if effect_outputs and spec.effect_mode != "none":
+        outputs = builder.emit_output(rx.Tuple([_unwrap_ret(outputs), rx.Tuple(effect_outputs)]))
+    else:
+        outputs = builder.emit_output(_unwrap_ret(outputs))
+    return outputs, inputs
+
+
+def _method_spec_to_inputs(
+    spec: _spec.MethodSpec,
+) -> list[tirx.Var | core.Tensor]:
+    """Convert the MethodSpec to a list of inputs to Module's method."""
+    str2var: dict[str, tirx.Var] = {}
+
+    def _get_var(name: str) -> tirx.Var:
+        if name in str2var:
+            return str2var[name]
+        var = tirx.Var(name, "int64")
+        str2var[name] = var
+        return var
+
+    def _convert_input(arg_name, arg_spec):
+        if isinstance(arg_spec, _spec.Int):
+            arg = _get_var(arg_name)
+        elif isinstance(arg_spec, _spec.Tensor):
+            arg = core.Tensor.placeholder(  # pylint: disable=protected-access
+                shape=[_get_var(x) if isinstance(x, str) else x for x in arg_spec.shape],
+                dtype=arg_spec.dtype,
+                name=arg_name,
+            )
+        elif isinstance(arg_spec, _spec.Object):
+            arg = arg_spec.object_type(_expr=rx.Var(arg_name, AnyType()), _name=arg_name)
+        elif isinstance(arg_spec, _spec.Tuple):
+            elements = type(arg_spec.elements)(
+                [
+                    _convert_input(arg_name=arg_name + f"_{i}", arg_spec=arg_spec.elements[i])
+                    for i in range(len(arg_spec.elements))
+                ]
+            )
+            arg = _spec.Tuple(
+                name=arg_name,
+                elements=elements,
+            )
+        else:
+            raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")
+        return arg
+
+    args = []
+    for arg_name, arg_spec in zip(spec.arg_names, spec.arg_specs):
+        arg = _convert_input(arg_name=arg_name, arg_spec=arg_spec)
+        args.append(arg)
+    return args

python/tvm/relax/frontend/nn/extern.py

@@ -0,0 +1,402 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E722
+"""External modules to be linked into the exported IRModule."""
+
+import os
+import shutil
+import sys
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+
+import tvm_ffi
+
+import tvm
+from tvm import libinfo, tirx
+from tvm.runtime import Module, load_static_library
+from tvm.support import cc as _cc
+
+from ...op import call_dps_packed
+from . import core
+from .core import wrap_nested
+
+
+class ExternModule:
+    """The abstract base class for external modules. External modules are designed to help
+    incorporate user-provided handcrafted kernels into the exported TVM IRModule.
+    """
+
+    _symbols: dict[str, Callable]
+
+    def __init__(self, symbols: dict[str, Callable]) -> None:
+        self._symbols = symbols
+
+    def __getitem__(self, func_name: str) -> Callable:
+        _inference_function = self._symbols[func_name]
+
+        def _call(*input_args):
+            def _convert(arg, name: str):
+                from tvm import relax as rx  # pylint: disable=import-outside-toplevel
+
+                if isinstance(arg, core.Tensor):
+                    return arg._expr  # pylint: disable=protected-access
+                if isinstance(arg, int):
+                    return rx.prim_value(tirx.IntImm("int64", arg))
+                if isinstance(arg, float):
+                    return rx.prim_value(tirx.FloatImm("float64", arg))
+                if isinstance(arg, str):
+                    return rx.StringImm(arg)
+                if tvm.ir.is_prim_expr(arg):
+                    return rx.prim_value(arg)
+                if isinstance(arg, tuple | list):
+                    return rx.Tuple([_convert(e, f"{name}_{i}") for i, e in enumerate(arg)])
+                raise TypeError(f"Unsupported input type: {type(arg)}")
+
+            rx_inputs = _convert(input_args, "input")
+            rx_outputs_ty = _convert(_inference_function(*input_args), "dummy").ty
+            return wrap_nested(call_dps_packed(func_name, rx_inputs, rx_outputs_ty), func_name)
+
+        return _call
+
+    def _load(self, path: Path) -> Module:
+        return load_static_library(str(path), func_names=list(self._symbols.keys()))
+
+    def load(self) -> Module:
+        """Loads the external module into a TVM runtime module."""
+        raise NotImplementedError
+
+
+class ObjectModule(ExternModule):  # pylint: disable=too-few-public-methods
+    """A subclass of `nn.ExternModule`, which allows
+    users to provide an object `.o` file to be linked into compiled
+    artifact;
+    """
+
+    def __init__(
+        self,
+        symbols: dict[str, Callable],
+        filepath: Path,
+    ) -> None:
+        if not isinstance(filepath, Path):
+            filepath = Path(filepath)
+        if not filepath.is_file():
+            raise ValueError(f"Not a file: {filepath!s}")
+        self.filepath = filepath
+        super().__init__(symbols)
+
+    def load(self) -> Module:
+        return self._load(self.filepath)
+
+
+class SourceModule(ExternModule):  # pylint: disable=too-few-public-methods
+    """A subclass of `nn.ExternModule`. It compiles C++/CUDA source code and link them into the
+    eventual IRModule.
+
+    **Shape/dtype inference.** The `nn.ExternModule` system requires users to provide additional
+    information to work, namely, `symbols`. It is a dictionary that maps each symbol in the
+    external object file to its shape/dtype inference function. Consider a case where function
+    `my_func` accepts two tensors, `a` of shape `(x, y, 1)`, and `b` of shape `(y, z, 5)`, and
+    produces a tensor `c` of shape `(x, y, z, 9)`, the shape/dtype inference function should look
+    like:
+
+    .. code-block:: python
+
+        def shape_dtype_inference(a, b):
+            x, y, _ = a.shape
+            _, z, _ = b.shape
+            return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")
+
+
+    and the `symbols` dictionary should be provided as:
+
+    .. code-block:: python
+
+        symbols={
+            "my_func": shape_dtype_inference,
+        }
+
+
+    **Calling convention.** All external modules now follows "destination-passing-style" (DPS)
+    calling convention, which means the returned tensors are pre-allocated by the system already
+    and passed in as an argument of the external function.
+
+    Reuse the example above, the implementation of `my_func` should include three parameters in
+    its signature, where tensors are represented using DLTensor from DLPack, the de facto standard
+    of in-memory representation of tensors. More details:
+    https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206.
+
+    To expose the symbol, `TVM_FFI_DLL_EXPORT_TYPED_FUNC(symbol, function)` is guaranteed available:
+
+    .. code-block:: C++
+
+        // those headers are guaranteed to be available
+        #include <dlpack/dlpack.h>
+        #include <tvm/ffi/dtype.h>
+        #include <tvm/ffi/function.h>
+
+        namespace {
+        // anonymous namespace hides the symbol `_my_func_impl` from other translation units
+        int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
+            // `a` and `b` are inputs, and `c` is the output
+        }
+        }
+        // expose symbol `my_func` instead of `_my_func_impl`
+        TVM_FFI_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);
+
+    **A compiler pass `AttachExternModules`.** It is introduced to attach a list of
+    `nn.ExternModule`s into an IRModule at any stage of the compilation pipeline,
+    and attach the compiled external modules as `runtime.Module`s into IRModule's `external_mods`
+    attribute. It is required by linking in `tvm.compile`, but with the existence of this pass,
+    source compilation can be deferred to arbitrary stage of TVM compilation.
+
+    **Caveats.** It is required to call `nn.add_extern` to register external modules exactly once
+    during `export_tvm`. Each symbol should be registered exactly once to avoid potential conflicts,
+    and otherwise an error will be raised.
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        symbols: dict[str, Callable],
+        source_code: str | Path,
+        source_format: str,  # "cpp", "cu"
+        compile_options: list[str] | None = None,
+        compiler: str | None = None,
+        output_format: str = "obj",  # "obj", "wasm"
+    ):
+        """Constructs a `nn.SourceModule` from source code.
+
+        Parameters
+        ----------
+        symbols : Dict[str, Callable]
+            The dictionary that maps each symbol in the external object file to its shape/dtype
+            inference function.
+
+        source_code : Union[str, Path]
+            Source code or path to the source code to be compiled.
+
+        source_format : str
+            The source code format. It can be either "cpp" or "cu".
+
+        compile_options : Optional[List[str]]
+            The compile options. If not provided, the default compile options will be used.
+
+        compiler : Optional[str]
+            The compiler. If not provided, the default compiler will be used. On Windows,
+            compilation requires `clang` by default.
+
+        output_format : str
+            The output format. It can be either "obj" or "wasm". "obj" is the default format,
+            which is a shared object file. "wasm" is the WebAssembly format, which is a binary
+            file.
+        """
+
+        def _detect_input_suffix(source_format: str) -> str:
+            if source_format == "cpp":
+                return ".cpp"
+            if source_format == "cu":
+                return ".cu"
+            raise ValueError(f"Invalid source format: {source_format}")
+
+        def _detect_output_suffix(output_format: str) -> str:
+            if output_format == "obj":
+                if _cc._is_linux_like():  # pylint: disable=protected-access
+                    return ".o"
+                if _cc._is_windows_like():  # pylint: disable=protected-access
+                    return ".obj"
+                raise ValueError(f"Unsupported platform: {sys.platform}")
+            if output_format == "wasm":
+                return ".wasm"
+            raise ValueError(f"Invalid output format: {output_format}")
+
+        def _detect_source_code(source_code) -> str:
+            if isinstance(source_code, Path):
+                path = source_code
+                if not path.is_file():
+                    raise ValueError(f"Not a file: {path!s}")
+            else:
+                try:
+                    path = Path(source_code)
+                except:  # pylint: disable=bare-except
+                    return source_code
+                try:
+                    if not path.is_file():
+                        return source_code
+                except:  # pylint: disable=bare-except
+                    return source_code
+            with path.open("r", encoding="utf-8") as file:
+                return file.read()
+
+        self.source_code = _detect_source_code(source_code)
+        if compile_options is None:
+            self.compile_options = SourceModule.get_compile_options(source_format=source_format)
+        else:
+            self.compile_options = list(compile_options)
+        self.compiler = compiler
+        self.source_suffix = _detect_input_suffix(source_format)
+        self.output_suffix = _detect_output_suffix(output_format)
+        super().__init__(symbols)
+
+    @staticmethod
+    def tvm_home() -> Path:
+        """Find TVM's home directory. If `TVM_HOME` environment variable is set, use it.
+        Otherwise, use the directory where the `tvm` Python package is installed.
+        As a sanity check, it is required to have `include` and `3rdparty` as direct subdirectories.
+
+        Returns
+        -------
+        tvm_home : pathlib.Path
+            The TVM home directory, and it is guaranteed to have `include` and `3rdparty` as
+            direct subdirectories.
+        """
+        if os.environ.get("TVM_HOME", None):
+            tvm_path = Path(os.environ["TVM_HOME"])
+            assert tvm_path.exists(), (
+                f"Using environment variable `TVM_HOME`, but directory not found: {tvm_path!s}"
+            )
+            assert tvm_path.is_dir(), (
+                f"Using environment variable `TVM_HOME`, but it is not a directory: {tvm_path!s}"
+            )
+        else:
+            import tvm  # pylint: disable=import-outside-toplevel
+
+            tvm_path = Path(tvm.__file__).parent
+            assert tvm_path.is_dir()
+        tvm_path = tvm_path.resolve()
+        while True:
+            exists_include = (tvm_path / "include").is_dir()
+            exists_3rdparty = (tvm_path / "3rdparty").is_dir()
+            if exists_include and exists_3rdparty:
+                return tvm_path.resolve()
+            parent = tvm_path.parent
+            if parent == tvm_path:
+                raise ValueError(
+                    "Cannot detect TVM directory. "
+                    "Please explicitly specify it by setting `TVM_HOME` environment variable, "
+                    "and make sure it contains `include` and `3rdparty` as direct sub-directories."
+                )
+            tvm_path = parent
+        return tvm_path.resolve()
+
+    @staticmethod
+    def get_includes(tvm_pkg: list[str] | None = None) -> list[Path]:
+        """Returns the default include paths according to `tvm_home()`.
+        By default, it includes TVM, DLPack. With `tvm_pkg` provided, it also
+        includes the specified package under `tvm_home/3rdparty`.
+
+        Parameters
+        ----------
+        tvm_pkg : Optional[List[str]]
+            The list of packages to be included under `tvm_home/3rdparty`. Each element should be
+            a relative path to `tvm_home/3rdparty`.
+
+        Returns
+        -------
+        includes : List[pathlib.Path]
+            The list of include paths.
+        """
+        results = [
+            Path(libinfo.find_include_path()),
+            Path(tvm_ffi.libinfo.find_include_path()),
+            Path(tvm_ffi.libinfo.find_dlpack_include_path()),
+        ]
+        if tvm_pkg:
+            tvm_home = SourceModule.tvm_home()
+            for relative in tvm_pkg:
+                results.append(tvm_home / "3rdparty" / relative)
+        results = list(dict.fromkeys(results))
+        for path in results:
+            assert path.exists(), f"Not found: {path!s}"
+            assert path.is_dir(), f"Not a directory: {path!s}"
+        return results
+
+    @staticmethod
+    def get_compile_options(
+        source_format: str,
+        tvm_pkg: list[str] | None = None,
+    ) -> list[str]:
+        """Returns the default compile options depending on `source_format`, including the default
+        inlcude paths w.r.t. `tvm_home()`, and by default,
+        it uses "-O3" and "-std=c++17".
+
+        Parameters
+        ----------
+        source_format : str
+            The source code format. It can be either "cpp" or "cu".
+
+        tvm_pkg : Optional[List[str]]
+            The list of packages to be included under `tvm_home/3rdparty`. Each element should be
+            a relative path to `tvm_home/3rdparty`.
+
+        Returns
+        -------
+        compile_options : List[str]
+            The list of compilation flags.
+        """
+        include_flags = []
+        for include_path in SourceModule.get_includes(tvm_pkg=tvm_pkg):
+            include_flags += ["-I", str(include_path)]
+        if source_format == "cpp":
+            host_flags = [
+                "-c",  # generate object file
+                "-O3",
+                "-std=c++17",
+            ]
+        elif source_format == "cu":
+            host_flags = [
+                "-c",  # generate object file
+                "-O3",
+                "-std=c++17",
+                # Enable `-fPIC` for the host compiler
+                "-Xcompiler=-fPIC",
+            ]
+        else:
+            raise ValueError(f"Invalid source format: {source_format}")
+        return include_flags + host_flags
+
+    def compile(self, output_path: Path) -> None:
+        """Compiles the source code in a provided directory and returns the compiled artifact."""
+        with tempfile.TemporaryDirectory() as temp_dir_str:
+            temp_dir = Path(temp_dir_str)
+            source_filename = f"main{self.source_suffix}"
+            object_filename = f"main{self.output_suffix}"
+            source_path = temp_dir / source_filename
+            object_path = temp_dir / object_filename
+            with source_path.open("w", encoding="utf-8") as file:
+                file.write(self.source_code)
+            _cc.create_shared(
+                output=object_filename,
+                objects=[source_filename],
+                options=self.compile_options,
+                cc=self.compiler,
+                cwd=temp_dir,
+                ccache_env=(
+                    {
+                        "CCACHE_COMPILERCHECK": "content",
+                        "CCACHE_NOHASHDIR": "1",
+                    }
+                    if shutil.which("ccache")
+                    else None
+                ),
+            )
+            shutil.move(str(object_path), str(output_path))
+
+    def load(self) -> Module:
+        with tempfile.TemporaryDirectory() as temp_dir_str:
+            output_path = Path(temp_dir_str) / f"main{self.output_suffix}"
+            self.compile(output_path)
+            return self._load(output_path)

python/tvm/relax/frontend/nn/llm/__init__.py

@@ -0,0 +1,23 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""LLM support for PyTorch-like API to build IRModules."""
+
+from . import kv_cache, position_embedding
+from .position_embedding import llama_rope
+from .tree_attn import tree_attn
+from .kv_cache import PagedKVCache

python/tvm/relax/frontend/nn/llm/_decode_kernels.py

@@ -0,0 +1,526 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501
+# fmt: off
+
+"""Single-token decode attention kernels and attention-state merge helpers.
+
+Contents:
+- ``_attention_decode_cpu`` / ``_attention_decode`` — paged-KV decode (one Q token
+  per sequence), CPU scalar and GPU allreduce variants.
+- ``_merge_state_inplace_cpu`` / ``_merge_state_inplace`` — combine two
+  log-sum-exp attention outputs in place. Used by multi-stage decoding and by
+  the distributed KV-transfer path.
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import math
+from typing import Any
+
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import (
+    _declare_length_info,
+    _get_kv_chunk_len,
+    _get_seq_offset,
+    _rope,
+    _var,
+    _var_cpu,
+    check_thread_limits,
+    get_max_num_threads_per_block,
+)
+
+
+def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], page_size: int = 16):
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+    group_size = num_qo_heads // num_kv_heads
+
+    global_symbol = "batch_decode_paged_kv_cpu"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle,  # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        # The length information of the sequences.
+        # - It is in shape `(3, batch_size)` when sliding window is enabled.
+        #   For a sequence "i", location
+        #   - "(0, i)" is the number of KV slots used in the last page of the seq ("last_page_len"),
+        #   - "(1, i)" is the starting offset of the sliding window in the seq,
+        #   - "(2, i)" is the attn sink length of the sequence.
+        # - It is in shape `(batch_size,)` when sliding window is disabled,
+        #   denoting the "last_page_len".
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for b in T.serial(B):
+            with T.sblock("attn"):
+                O_local = T.sblock_alloc_buffer((D,), "float32")
+                Q_local = T.sblock_alloc_buffer((D,), "float32")
+                K_local = T.sblock_alloc_buffer((D,), "float32")
+                V_local = T.sblock_alloc_buffer((D,), "float32")
+
+                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32")
+
+                m_val = T.sblock_alloc_buffer((1,), "float32")
+                new_m = T.sblock_alloc_buffer((1,), "float32")
+                d_val = T.sblock_alloc_buffer((1,), "float32")
+                S_val = T.sblock_alloc_buffer((1,), "float32")
+                scale_O = T.sblock_alloc_buffer((1,), "float32")
+                factor = T.sblock_alloc_buffer((1,), "float32")
+
+                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[b]
+                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[b + 1]
+
+                kv_chunk_len[0] = T.if_then_else(
+                    cur_page_indptr_begin != cur_page_indptr_end,
+                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b, length_info, sliding_window),
+                    0,
+                )
+
+                for h_qo in T.serial(H_qo):
+                    m_val[0] = -5e4
+                    d_val[0] = 1.0
+
+                    for d in T.serial(D):
+                        O_local[d] = 0.0
+
+                    for d in T.serial(D):
+                        Q_local[d] = T.if_then_else(
+                            rotary_mode == 1,
+                            _rope(Q, q_rope_position[b], head_dim, rope_theta, rope_scale, (b, h_qo, d), qkv_dtype, rope_scaling),
+                            Q[b, h_qo, d],
+                        )
+
+                    for row_idx in T.serial(kv_chunk_len[0]):
+                        seq_offset: T.let[T.int32()] = _get_seq_offset(row_idx, b, length_info, sliding_window)
+                        page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
+                        page_offset: T.let[T.int32()] = seq_offset % page_size
+
+                        for d in T.serial(D):
+                            K_local[d] = T.if_then_else(
+                                rotary_mode == 1,
+                                _rope(pages, k_rope_pos_offset[b] + row_idx, head_dim, rope_theta, rope_scale, (page_no, 0, h_qo // group_size, page_offset, d), qkv_dtype, rope_scaling),
+                                pages[page_no, 0, h_qo // group_size, page_offset, d],
+                            )
+                        S_val[0] = 0.0
+                        for d in T.serial(D):
+                            S_val[0] += Q_local[d] * K_local[d]
+                        S_val[0] *= sm_scale * math.log2(math.exp(1))
+
+                        new_m[0] = T.max(m_val[0], S_val[0])
+                        d_val[0] = (d_val[0] * T.exp2(m_val[0] - new_m[0])) + T.exp2(S_val[0] - new_m[0])
+
+                        scale_O[0] = T.exp2(m_val[0] - new_m[0])
+
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] * scale_O[0]
+
+                        m_val[0] = new_m[0]
+                        for d in T.serial(D):
+                            V_local[d] = pages[page_no, 1, h_qo // group_size, page_offset, d]
+
+                        factor[0] = T.exp2(S_val[0] - m_val[0])
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] + V_local[d] * factor[0]
+                    for d in T.serial(D):
+                        O_local[d] = O_local[d] / d_val[0]
+                        output[b, h_qo, d] = O_local[d]
+                    lse[b, h_qo] = m_val[0] + T.log2(d_val[0])
+
+    return batch_decode_paged_kv
+
+
+def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16):
+    qkv_dtype_bytes = 2
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+
+    THREAD_LIMIT = 512
+    TILE_SIZE_PER_BDX = 2
+    if target.kind.name == "opencl" and (("android" in str(target.host)) or ("adreno" in str(target.attrs))):
+        # Keeping lower thread limit for this kernel on adreno target
+        # to avoid register spill
+        THREAD_LIMIT = 256
+        TILE_SIZE_PER_BDX = 1
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    thread_limit = min(max_num_threads_per_block, THREAD_LIMIT)
+
+    GROUP_SIZE = H_qo // H_kv
+    VEC_SIZE = min(max(8 // qkv_dtype_bytes, D // 32), 4)
+    bdx = D // VEC_SIZE
+    bdy = GROUP_SIZE
+    while bdx * bdy > thread_limit and bdy > 1:
+        bdy //= 2
+    gdz = GROUP_SIZE // bdy
+    threads_per_CTA = max(thread_limit, bdx * bdy)
+    bdz = threads_per_CTA // (bdx * bdy)
+    tile_size_per_bdx = TILE_SIZE_PER_BDX if GROUP_SIZE == 1 else 1
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=bdz, gdz=1)
+
+    global_symbol = "batch_decode_paged_kv"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    # pylint: disable=too-many-branches
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for bx in T.thread_binding(B, thread="blockIdx.x"):
+            for fused_by_bz in T.thread_binding(H_kv * gdz, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        for tz in T.thread_binding(bdz, thread="threadIdx.z"):
+                            with T.sblock("attn"):
+                                Q_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32", scope="local")
+                                K_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                V_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                O_allreduce = T.sblock_alloc_buffer((bdz, bdy, D), "float32", scope="shared")
+                                md_allreduce = T.sblock_alloc_buffer((bdz, bdy, 2), "float32", scope="shared")
+                                S_reduce_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
+
+                                S_local = T.sblock_alloc_buffer((bdy * tile_size_per_bdx), "float32", scope="local")
+                                QK_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                V_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                m_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                d_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_mprev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_otherm = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_o = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                st_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                st_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                O_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+
+                                by: T.let[T.int32] = fused_by_bz % H_kv
+                                bz: T.let[T.int32] = fused_by_bz // H_kv
+                                batch_idx: T.let[T.int32] = bx
+                                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[batch_idx]
+                                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[batch_idx + 1]
+                                kv_chunk_len[0] = T.if_then_else(
+                                    cur_page_indptr_begin != cur_page_indptr_end,
+                                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, batch_idx, length_info, sliding_window),
+                                    0
+                                )
+
+                                # init states
+                                st_m[0] = -5e4
+                                st_d[0] = 1.0
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] = 0.0
+
+                                # load q
+                                for vec in T.vectorized(VEC_SIZE):
+                                    Q_local[vec] = T.if_then_else(
+                                        rotary_mode == 1,
+                                        _rope(Q, q_rope_position[batch_idx], head_dim, rope_theta, rope_scale, (bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                        Q[bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec]
+                                    )
+
+                                for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_size_per_bdx * bdy * bdz)):
+                                    tile_start_s: T.let[T.int32()] = (tz * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    tile_start_g: T.let[T.int32()] = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    # load KV from global memory to shared memory
+                                    for j in T.serial(tile_size_per_bdx):
+                                        with T.sblock("KV_load"):
+                                            T.reads()
+                                            T.writes()
+                                            row_g: T.let[T.int32()] = tile_start_g + j  # type: ignore
+                                            if row_g < kv_chunk_len[0]:
+                                                seq_offset: T.let[T.int32()] = _get_seq_offset(row_g, batch_idx, length_info, sliding_window)  # type: ignore
+                                                page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)]  # type: ignore
+                                                page_offset: T.let[T.int32()] = T.floormod(seq_offset, page_size)  # type: ignore
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = T.if_then_else(
+                                                        rotary_mode == 1,
+                                                        _rope(pages, k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale, (page_no, 0, by, page_offset, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                                        pages[page_no, 0, by, page_offset, tx * VEC_SIZE + vec]
+                                                    )
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = pages[page_no, 1, by, page_offset, tx * VEC_SIZE + vec]
+                                            else:
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                    T.tvm_storage_sync("shared")
+                                    # compute QK
+                                    m_prev[0] = st_m[0]
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        # compute S = Q * K * sm_scale
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            QK_local[vec] = T.cast(Q_local[vec], "float32") * T.cast(K_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec], "float32") * sm_scale * math.log2(math.exp(1))
+                                        S_reduce_local[0] = 0
+                                        for vec in T.unroll(VEC_SIZE):
+                                            S_reduce_local[0] += QK_local[vec]
+
+                                        with T.sblock("block_cross_thread"):
+                                            T.reads(S_reduce_local[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), S_reduce_local[0], True, t0[0], tx, dtype="void")
+
+                                        S_local[j] = -5e4
+                                        if (iterator * bdz + tz) * bdy * tile_size_per_bdx + j < kv_chunk_len[0]:
+                                            S_local[j] = t0[0]
+                                        # update st_m
+                                        st_m[0] = T.max(st_m[0], S_local[j])
+
+                                    # update st_d, st_O
+                                    o_scale: T.let[T.float32] = T.exp2(m_prev[0] - st_m[0])
+                                    st_d[0] *= o_scale
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        S_local[j] = T.exp2(S_local[j] - st_m[0])
+                                        st_d[0] += S_local[j]
+                                    for j in T.vectorized(VEC_SIZE):
+                                        O_local[j] *= o_scale
+
+                                    # load V from shared memory to local memory
+                                    # compute O
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            V_local[vec] = V_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] += T.cast(V_local[vec], "float32") * S_local[j]
+
+                                if bdz > 1:
+                                    # allreduce over bdz
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_allreduce[tz, ty, tx * VEC_SIZE + vec] = O_local[vec]
+                                    md_allreduce[tz, ty, 0] = st_m[0]
+                                    md_allreduce[tz, ty, 1] = st_d[0]
+                                    T.tvm_storage_sync("shared")
+
+                                    st_m[0] = -5e4
+                                    st_d[0] = 1.0
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_local[vec] = 0.0
+
+                                    for j in T.serial(bdz):
+                                        m_prev[0] = st_m[0]
+                                        d_prev[0] = st_d[0]
+                                        other_m[0] = md_allreduce[j, ty, 0]
+                                        other_d[0] = md_allreduce[j, ty, 1]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            other_o[vec] = O_allreduce[j, ty, tx * VEC_SIZE + vec]
+                                        st_m[0] = T.max(st_m[0], other_m[0])
+                                        st_d[0] = d_prev[0] * T.exp2(m_prev[0] - st_m[0]) + other_d[0] * T.exp2(other_m[0] - st_m[0])
+                                        exp_mprev[0] = T.exp2(m_prev[0] - st_m[0])
+                                        exp_otherm[0] = T.exp2(other_m[0] - st_m[0])
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] = O_local[vec] * exp_mprev[0] + other_o[vec] * exp_otherm[0]
+
+                                # normalize O
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] /= st_d[0]
+
+                                # store O to global memory
+                                for vec in T.vectorized(VEC_SIZE):
+                                    output[batch_idx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec] = O_local[vec]
+
+                                # store lse to global memory
+                                lse[batch_idx, by * GROUP_SIZE + bz * bdy + ty] = st_m[0] + T.log2(st_d[0])
+    # pylint: enable=too-many-branches
+    return batch_decode_paged_kv
+
+
+def _merge_state_inplace_cpu(v_dtype):
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace_cpu(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for n in T.serial(N):
+            for h in T.serial(H):
+                with T.sblock("merge"):
+                    s_val = _var_cpu("float32")
+                    s_other_val = _var_cpu("float32")
+                    s_max = _var_cpu("float32")
+                    scale = _var_cpu("float32")
+                    other_scale = _var_cpu("float32")
+
+                    s_val[0] = S[n, h]
+                    s_other_val[0] = S_other[n, h]
+                    s_max[0] = T.max(s_val[0], s_other_val[0])
+                    s_val[0] = T.exp2(s_val[0] - s_max[0])
+                    s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                    scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                    other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+                    for d in T.serial(D):
+                        V[n, h, d] = V[n, h, d] * scale[0] + V_other[n, h, d] * other_scale[0]
+                    S[n, h] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    return merge_state_inplace_cpu
+
+
+def _merge_state_inplace(num_heads, head_dim, v_dtype, target: Target, global_symbol: str | None = None):
+    v_dtype_bytes = 2
+    VEC_SIZE = min(max(8 // v_dtype_bytes, head_dim // 32), 4)
+    bdx = head_dim // VEC_SIZE
+    bdy = num_heads
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    while bdx * bdy > max_num_threads_per_block and bdy > 1:
+        bdy //= 2
+    gdy = num_heads // bdy
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for bx in T.thread_binding(N, thread="blockIdx.x"):
+            for by in T.thread_binding(gdy, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        with T.sblock("merge"):
+                            s_val = _var("float32")
+                            s_other_val = _var("float32")
+                            s_max = _var("float32")
+                            scale = _var("float32")
+                            other_scale = _var("float32")
+
+                            v_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+                            v_other_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+
+                            s_val[0] = S[bx, ty + by * bdy]
+                            s_other_val[0] = S_other[bx, ty + by * bdy]
+                            s_max[0] = T.max(s_val[0], s_other_val[0])
+                            s_val[0] = T.exp2(s_val[0] - s_max[0])
+                            s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                            scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                            other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+
+                            # load v
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_vec[vec] = V[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+                            # load v_other
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_other_vec[vec] = V_other[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+
+                            # merge
+                            for vec in T.serial(VEC_SIZE):
+                                v_vec[vec] = v_vec[vec] * scale[0] + v_other_vec[vec] * other_scale[0]
+
+                            # store v
+                            for vec in T.vectorized(VEC_SIZE):
+                                V[bx, ty + by * bdy, tx * VEC_SIZE + vec] = v_vec[vec]
+
+                            # store s
+                            S[bx, ty + by * bdy] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    func = merge_state_inplace
+    if global_symbol:
+        func = func.with_attr("global_symbol", global_symbol)
+    return func

python/tvm/relax/frontend/nn/llm/_kernel_common.py

@@ -0,0 +1,569 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, E731, RUF005
+# fmt: off
+
+"""Shared TIR helpers used by KV-cache / attention kernels in this package.
+
+This module consolidates constructs reused by the prefill/decode/paged/tree
+attention kernels so each kernel file can focus on its own specialised logic.
+
+Contents:
+- Thread-limit checks (``get_max_num_threads_per_block``, ``check_thread_limits``)
+- KV-cache enums (``AttnKind``, ``RopeMode``)
+- Small TVMScript helpers (``_var``, ``_var_cpu``, ``_causal_mask``, ``_rope``)
+- Length-info accessors for sliding-window-aware indexing
+- Buffer allocators for the tiled online-softmax state used by every prefill kernel
+- ``_make_prefill_macros`` — the ``@T.macro`` bundle invoked by the prefill kernels
+- Tiling config (``_get_prefill_kernel_config``) and scheduling (``_schedule_prefill_kernel``)
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import enum
+import math
+from typing import Any
+
+import tvm
+from tvm import s_tir, tirx
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from .position_embedding import switch_rope_freq_func
+
+
+def _var(dtype):
+    return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+
+def _var_cpu(dtype):
+    return T.sblock_alloc_buffer((1,), dtype)
+
+
+def get_max_num_threads_per_block(target: Target) -> int:
+    """
+    max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads.
+    We add this method since some targets have both fields and `max_threads_per_block` is larger.
+    """
+    max_num_threads = int(target.attrs["max_num_threads"])
+    max_threads_per_block = target.attrs.get("max_threads_per_block", None)
+    if max_threads_per_block is None:
+        return max_num_threads
+    return max(max_num_threads, max_threads_per_block)
+
+
+def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int):
+    """
+    Check whether max num threads exceeded given a target.
+
+    Parameters
+    ----------
+    bdx: threadIdx.x
+    bdy: threadIdx.y
+    bdz: threadIdx.z
+    gdz: blockIdx.z
+    """
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+
+    assert bdx * bdy * bdz <= max_num_threads_per_block, (
+        f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}"
+    )
+
+    if target.kind.name == "webgpu":
+        # https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez
+        assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}"
+        assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}"
+
+
+class AttnKind(enum.IntEnum):
+    """The attention kind class.
+    MHA denotes multi-head attention, multi-query attention or grouped query attention.
+    MLA denotes multi-head latent attention.
+    """
+
+    MHA = 0
+    MLA = 1
+    MHA_SLIDING = 3
+
+
+class RopeMode(enum.IntEnum):
+    """The RoPE mode of the Paged KV cache.
+    If it is none, the KV cache will not apply RoPE to q and k.
+    If it is normal, RoPE will be applied to k before adding k to cache.
+    Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+    """
+
+    NONE = 0
+    NORMAL = 1
+    INLINE = 2
+
+
+def _rope(buffer: T.Buffer, offset: tirx.Var, rotary_dim: int, theta: tirx.Var, scale: tirx.Var, indices: tuple[tirx.Var, ...], qkv_dtype: str, rope_scaling: dict[str, Any]):
+    d = indices[-1]
+    cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(offset * scale, d, rotary_dim, theta, "float32")
+    cos = cos_freq * buffer[indices].astype("float32")
+    sin = sin_freq * tirx.if_then_else(
+        d < rotary_dim // 2,
+        -buffer[indices[:-1] + (d + rotary_dim // 2,)],
+        buffer[indices[:-1] + (d - rotary_dim // 2,)],
+    ).astype("float32")
+    expr = (cos + sin).astype(qkv_dtype)
+    for var, value in var_map.items():
+        expr = tirx.Let(var, value, expr)
+    return expr
+
+
+def _causal_mask(causal, row, col, kv_len, qo_len):
+    return T.if_then_else(
+        causal > 0,
+        col < kv_len - qo_len + row + 1,
+        col < kv_len,
+    )
+
+
+def _declare_length_info(var_length_info, batch_size, sliding_window, elem_offset):
+    return (
+        T.match_buffer(var_length_info, (3, batch_size), "int32", elem_offset=elem_offset)
+        if sliding_window
+        else T.match_buffer(var_length_info, (batch_size,), "int32", elem_offset=elem_offset)
+    )
+
+
+def _get_kv_chunk_len(num_pages, page_size, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return (num_pages - 1) * page_size + length_info[seq_id]
+    # ((num_pages - 1) * page_size + last_page_len) - sliding_window_offset + sink_size
+    return (num_pages - 1) * page_size + length_info[0, seq_id] - length_info[1, seq_id] + length_info[2, seq_id]
+
+
+def _get_seq_offset(pos, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return pos
+    # pos if pos < sink_size else pos - sink_size + sliding_window_offset
+    return T.if_then_else(
+        pos < length_info[2, seq_id],
+        pos,
+        pos - length_info[2, seq_id] + length_info[1, seq_id],
+    )
+
+
+def _alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps):
+    """Allocate the shared/local online-softmax working state used by every tiled prefill kernel.
+
+    Returns ``(S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new)``.
+    """
+    S_smem = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="shared")
+    S_local = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="local")
+    m_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    m_prev_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    d_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    md_shape = (math.ceil(tile_x / (bdx * num_warps)),)
+    m_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    m_prev = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    d_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    return S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new
+
+
+def _alloc_mha_qkvo_buffers(tile_x, tile_z, d_qk, d_v, dtype):
+    """Allocate Q/K/V shared + O local buffers for standard MHA/GQA prefill kernels."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    K_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    V_smem = T.sblock_alloc_buffer((tile_z, d_v), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_v), "float32", scope="local")
+    return Q_smem, K_smem, V_smem, O_local
+
+
+def _alloc_mla_qkvo_buffers(tile_x, tile_z, d_qk, d_latent, dtype):
+    """Allocate Q + combined KV shared + O local for MLA prefill (V reuses the KV buffer)."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    KV_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_latent), "float32", scope="local")
+    return Q_smem, KV_smem, O_local
+
+
+def _alloc_tile_walk_state():
+    """Return (tile_id, batch_idx, batch_tiles, batch_rows, iterator, kv_chunk_len) int32 scalars for the paged/ragged/MLA tile-walk state machine."""
+    return _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32")
+
+
+def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_size):
+    """Build @T.macro helpers shared across tiled online-softmax prefill kernels.
+
+    Parameters
+    ----------
+    tile_x : int  # query/output row tile
+    tile_y : int  # QK reduction dim (head_dim for MHA, d_qk for MLA/ragged)
+    tile_z : int  # key/value column tile
+    tile_o : int  # output/V column dim (d for MHA/sequence, d_v for ragged, d_latent for MLA)
+    """
+    @T.macro
+    def init_states(
+        m_smem: T.Buffer, d_smem: T.Buffer, O_local: T.Buffer, ty: T.int32, tx: T.int32,
+    ):
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                m_smem[row] = -5e4
+                d_smem[row] = 1.0
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_init"):
+                i, j = T.axis.remap("SS", [li, lj])
+                O_local[i, j] = 0.0
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_s_gemm(
+        Q_smem: T.Buffer, K_smem: T.Buffer, S_local: T.Buffer, S_smem: T.Buffer, sm_scale: T.float32,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_z, tile_y):
+                with T.sblock("S_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        S_local[i, j] = 0.0
+                    S_local[i, j] += T.cast(Q_smem[i, k], "float32") * T.cast(K_smem[j, k], "float32") * sm_scale * math.log2(math.exp(1))
+        T.tvm_storage_sync("shared")
+        for li, lj in T.grid(tile_x, tile_z):
+            with T.sblock("S_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                S_smem[i, j] = S_local[i, j]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        causal: T.int32, kv_len: T.int32, qo_len: T.int32,
+    ):
+        # Phase 1: compute m_new = max(masked S over kv tile), d_new = d_prev * exp2(m_prev - m_new)
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        # Phase 2: exp-and-scale S_smem; masked-out entries use -inf
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    # predicate sits inside loop so sync stays outside conditional branches
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        # Phase 3: d_new += sum(S_smem[row, :]); write m/d/m_prev back to smem
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_o_gemm(
+        S_smem: T.Buffer, V_smem: T.Buffer, O_local: T.Buffer,
+        m_prev_smem: T.Buffer, m_smem: T.Buffer,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_o, tile_z):
+                with T.sblock("O_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        O_local[i, j] *= T.exp2(m_prev_smem[i] - m_smem[i])
+                    O_local[i, j] += S_smem[i, k] * T.cast(V_smem[k, j], "float32")
+
+    @T.macro
+    def paged_store_output_lse(
+        output: T.Buffer, lse: T.Buffer, O_local: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer,
+        q_indptr: T.Buffer, b_idx: T.int32, by: T.int32, LH_start: T.int32,
+    ):
+        """Paged-style (q_indptr-based) O_store + lse_store epilogue.
+
+        Used by paged prefill, ragged prefill and MLA prefill. MLA passes ``by=0`` so
+        the ``by * group_size`` term drops to zero at compile time.
+        """
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    output[cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
+        for li in T.grid(tile_x):
+            with T.sblock("lse_store"):
+                i = T.axis.remap("S", [li])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    lse[cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
+
+    @T.macro
+    def advance_tile_batch(
+        tile_id: T.Buffer, batch_idx: T.Buffer, batch_tiles: T.Buffer, batch_rows: T.Buffer,
+        q_indptr: T.Buffer, batch_size: T.int32,
+    ):
+        """Advance tile_id/batch_idx past exhausted batches.
+
+        After the loop, either batch_idx[0] >= batch_size (all tiles consumed) or
+        tile_id[0] < batch_tiles[0] (the current batch still has work to do).
+        """
+        while tile_id[0] >= batch_tiles[0] and batch_idx[0] < batch_size:
+            tile_id[0] -= batch_tiles[0]
+            batch_idx[0] += 1
+            if batch_idx[0] < batch_size:
+                b_idx: T.let[T.int32] = batch_idx[0]
+                batch_rows[0] = (q_indptr[b_idx + 1] - q_indptr[b_idx]) * group_size
+                batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
+
+    @T.macro
+    def softmax_update_valid_length(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Same three-phase online softmax as softmax_update_causal but with a
+        # per-batch right-padding mask in place of causal masking.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal_padded_left(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Three-phase online softmax with left-padding + causal mask. Real
+        # queries occupy [qo_len - valid_len, qo_len); real keys occupy
+        # [kv_len - valid_len, kv_len). Causal keeps
+        # col <= row + (kv_len - qo_len) within those valid suffixes.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    pad_q: T.let[T.int32] = qo_len - valid_len
+                    pad_kv: T.let[T.int32] = kv_len - valid_len
+                    for j in T.serial(tile_z):
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        pad_q: T.let[T.int32] = qo_len - valid_len
+                        pad_kv: T.let[T.int32] = kv_len - valid_len
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    return init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, softmax_update_valid_length, advance_tile_batch, paged_store_output_lse, softmax_update_causal_padded_left
+
+
+def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target):
+    NUM_BLKS = 16
+    LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8)  # 8 bytes
+    group_size = h_q // h_kv
+
+    bdx = 32
+    num_warps = 4
+    tile_x, tile_y, tile_z = (
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+        d,
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+    )
+    original_tile_y = tile_y
+    original_tile_z = tile_z
+    while (tile_x * tile_z) % (bdx * num_warps) != 0:
+        tile_z += original_tile_z
+    while (tile_x * tile_y) % (bdx * num_warps) != 0:
+        tile_y += original_tile_y
+
+    # Otherwise we would exceed maxComputeWorkgroupStorageSize
+    if (
+        target.kind.name == "webgpu"
+        and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
+    ):
+        tile_z = 8
+        num_warps = 2
+    if target.kind.name == "opencl" and (
+        ("android" in str(target.host)) or ("adreno" in str(target.attrs))
+    ):
+        LOAD_VEC = 16 // ((DataType(dtype).bits + 7) // 8)  # 16 bytes
+        NUM_BLKS = group_size * 8
+
+    check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)
+
+    return NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z
+
+
+def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec, bdx, num_warps, tile_x, tile_y, tile_z, transform_k_load: bool, merged_qk_load: bool) -> tvm.s_tir.Schedule:
+    get_extent = lambda *lps: [int(sch.get(lp).extent) for lp in lps]
+
+    def get_vecsize(extent):
+        return min(load_vec, (extent & ~(extent - 1)))
+
+    def getxy_vecsize(x, y, t):
+        assert (x * y) % t == 0
+        return min(get_vecsize(y), get_vecsize(x * y // t))
+
+    def get_tile_size(x, y, t):
+        cnt = (x * y) // t
+        assert (x * y) % t == 0
+        tile_y = math.ceil(math.sqrt(cnt))
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 0) and tile_y <= cnt:
+            tile_y += 1
+        assert tile_y <= cnt
+        tile_x = cnt // tile_y
+        return tile_x, tile_y
+
+    def apply_to_qkv_load(sch: s_tir.Schedule, block):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        x_extent, y_extent = get_extent(loop_x, loop_y)
+        vec_size = getxy_vecsize(x_extent, y_extent, bdx * num_warps)
+        yo, yv = sch.split(loop_y, [None, vec_size])
+        yo_extent = y_extent // vec_size
+        tile_x, tile_y = get_tile_size(x_extent, yo_extent, (bdx * num_warps))
+        xo, xi = sch.split(loop_x, [tile_x, None])
+        yo, yi = sch.split(yo, [tile_y, None])
+        sch.reorder(xi, yi, xo, yo)
+        t = sch.fuse(xi, yi)
+        ty, tx = sch.split(t, [num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.vectorize(yv)
+
+    def apply_to_so_ewise(sch: s_tir.Schedule, block, tile):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    def apply_to_gemm(sch: s_tir.Schedule, block, tile, r_len=16, k_major=False):
+        loop_x, loop_y, loop_z = sch.get_loops(block)[-3:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+        ko, ki = sch.split(loop_z, factors=[None, r_len])
+        if k_major:
+            sch.reorder(ko, xi, yi, ki)
+        else:
+            sch.reorder(ko, ki, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        sch.unroll(xi)
+        sch.decompose_reduction(block, ty)
+
+    def apply_to_md(sch, block):
+        loop = sch.get_loops(block)[-1]
+        _, ty, tx = sch.split(loop, factors=[None, num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    if transform_k_load and not merged_qk_load:
+        sch.transform_layout("K_load", ("write", 0), lambda i, j: (j, i))
+    tile_s = get_tile_size(tile_x, tile_z, bdx * num_warps)
+    tile_o = get_tile_size(tile_x, tile_y, bdx * num_warps)
+    apply_to_gemm(sch, sch.get_sblock("S_gemm"), tile_s, k_major=True)
+    apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, k_major=False)
+    apply_to_so_ewise(sch, sch.get_sblock("S_store"), tile_s)
+    apply_to_so_ewise(sch, sch.get_sblock("O_init"), tile_o)
+    apply_to_so_ewise(sch, sch.get_sblock("O_store"), tile_o)
+    apply_to_qkv_load(sch, sch.get_sblock("Q_load"))
+    if not merged_qk_load:
+        apply_to_qkv_load(sch, sch.get_sblock("K_load"))
+        apply_to_qkv_load(sch, sch.get_sblock("V_load"))
+    else:
+        apply_to_qkv_load(sch, sch.get_sblock("KV_load"))
+    apply_to_md(sch, sch.get_sblock("lse_store"))
+    return sch

python/tvm/relax/frontend/nn/llm/_page_kernels.py

@@ -0,0 +1,293 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501
+# fmt: off
+
+"""TIR kernels that operate on paged KV-cache storage (without doing attention).
+
+This module contains:
+- Append helpers that transpose/write new K/V tokens into the paged layout
+  (``_kv_cache_transpose_append`` and its MLA variant).
+- Debug helpers that extract K/V from the paged layout for inspection
+  (``_kv_cache_debug_get_kv``, ``_kv_cache_debug_get_kv_mla``).
+- Copy helpers used by the cache runtime for forking/sharing pages
+  (``_copy_single_page``, ``_copy_single_page_mla``, ``_copy_single_page_cpu``).
+- Compact helpers that reorganise pages after removals
+  (``_compact_kv_copy``, ``_compact_kv_copy_cpu``).
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import get_max_num_threads_per_block
+
+
+def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size: int = 16):
+    """Return the TIR function that appends new k/v data to PagedKVCache."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append(
+        var_pages: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        k_data = T.match_buffer(var_k_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, h, f in T.grid(ntoken, num_key_value_heads, head_dim):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], k_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 0, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vf] = k_data[vgpos, vh, vf]
+                with T.sblock("v_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], v_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 1, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos] # type: ignore[name-defined,no-redef]
+                    pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vf] = v_data[vgpos, vh, vf]
+
+    return tir_kv_cache_transpose_append
+
+
+def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
+    """Return the TIR function that appends new compressed KV data to PagedKVCache for MLA."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append_mla(
+        var_pages: T.handle,
+        var_kv_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        kv_data = T.match_buffer(var_kv_data, (ntoken, d_qk), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, f in T.grid(ntoken, d_qk):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vf = T.axis.remap("SS", [global_pos, f])
+                    T.reads(position_map[vgpos], kv_data[vgpos, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), T.floormod(position, page_size), vf] = kv_data[vgpos, vf]
+
+    return tir_kv_cache_transpose_append_mla
+
+
+def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype,elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        k_data = T.match_buffer(var_k_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        for p, h, d in T.grid(seqlen, num_key_value_heads, head_dim):
+            with T.sblock("copy0"):
+                vp, vh, vd = T.axis.remap("SSS", [p, h, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, 0:2, vh, position_map[vp] % page_size, vd])
+                T.writes(k_data[layer_id, vp, vh, vd], v_data[layer_id, vp, vh, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                k_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vd]
+                v_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv
+
+
+def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv_mla(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_compressed_kv_with_k_pe_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        compressed_kv_with_k_pe_data = T.match_buffer(var_compressed_kv_with_k_pe_data, (num_hidden_layers, seqlen, d_qk), dtype)
+        for p, d in T.grid(seqlen, d_qk):
+            with T.sblock("copy0"):
+                vp, vd = T.axis.remap("SS", [p, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, position_map[vp] % page_size, vd])
+                T.writes(compressed_kv_with_k_pe_data[layer_id, vp, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                compressed_kv_with_k_pe_data[layer_id, vp, vd] = pages[T.floordiv(position, page_size), T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv_mla
+
+
+def _copy_single_page(num_heads, page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page
+
+
+def _copy_single_page_mla(page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_mla(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * head_dim)
+                    vp = T.axis.spatial(copy_length, (b * tx + t) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, vp, vd] = pages[src_page_id, vp, vd]
+
+    return copy_single_page_mla
+
+
+def _copy_single_page_cpu(num_heads, page_size, head_dim, dtype):
+    tx = 1
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_cpu(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+
+        for b in T.serial((copy_length * num_heads * head_dim + tx - 1) // tx):
+            for t in T.serial(tx):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page_cpu
+
+
+def _compact_kv_copy(num_heads, head_dim, dtype, target: Target, page_size: int = 16):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.thread_binding((batch_size * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+                for bhd_i in T.thread_binding(tx, thread="threadIdx.x"):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy
+
+
+def _compact_kv_copy_cpu(num_heads, head_dim, dtype, page_size: int = 16):
+    tx = 8
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy_cpu(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.serial((batch_size * num_heads * head_dim + tx - 1) // tx):
+                for bhd_i in T.serial(tx):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy_cpu

python/tvm/relax/frontend/nn/llm/kv_cache.py

@@ -0,0 +1,686 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E501, RUF012
+# fmt: off
+
+"""Attention KV cache modeling.
+
+This module exposes the public ``PagedKVCache`` classes (``FlashInferPagedKVCache``
+and ``TIRPagedKVCache``). The kernel factories that build the underlying TIR
+functions are split across sibling private modules:
+
+- ``_kernel_common``: shared helpers (enums, RoPE, mask, tile allocators,
+  ``@T.macro`` bundle, tiling config, scheduling).
+- ``_page_kernels``: page management (append, debug, copy, compact).
+- ``_prefill_kernels``: prefill attention kernels (paged/ragged/MLA/dense).
+- ``_decode_kernels``: decode attention kernels and state-merge helpers.
+
+The private-named kernel factories are re-exported from this module so the
+test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache``.
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import math
+from typing import Any, Literal
+
+import tvm
+from tvm import relax as rx
+from tvm import tirx
+from tvm.relax.frontend.nn import Object, Tensor
+from tvm.target import Target
+
+# Re-export enums + kernel factories so existing ``from kv_cache import ...``
+# users (test suite, tree_attn.py, mlc-llm, etc.) continue to work after the
+# split. These names are referenced in ``__all__`` below to signal to linters
+# that the imports are intentional public API (not dead code).
+from ._decode_kernels import (
+    _attention_decode,
+    _attention_decode_cpu,
+    _merge_state_inplace,
+    _merge_state_inplace_cpu,
+)
+from ._kernel_common import AttnKind, RopeMode
+from ._page_kernels import (
+    _compact_kv_copy,
+    _compact_kv_copy_cpu,
+    _copy_single_page,
+    _copy_single_page_cpu,
+    _copy_single_page_mla,
+    _kv_cache_debug_get_kv,
+    _kv_cache_debug_get_kv_mla,
+    _kv_cache_transpose_append,
+    _kv_cache_transpose_append_mla,
+)
+from ._prefill_kernels import (
+    _attention_prefill,
+    _attention_prefill_cpu,
+    _attention_prefill_mla,
+    _attention_prefill_ragged,
+    _attention_prefill_ragged_cpu,
+    _attention_sequence_prefill,
+    _attention_sequence_prefill_with_mask,
+)
+from .position_embedding import llama_rope_with_position_map
+from .tree_attn import (
+    tree_attn,
+    tree_attn_cpu,
+    tree_attn_with_paged_kv_cache,
+    tree_attn_with_paged_kv_cache_cpu,
+)
+
+__all__ = [
+    "AttnKind",
+    "FlashInferPagedKVCache",
+    "PagedKVCache",
+    "RopeMode",
+    "TIRPagedKVCache",
+    "_attention_decode",
+    "_attention_decode_cpu",
+    "_attention_prefill",
+    "_attention_prefill_cpu",
+    "_attention_prefill_mla",
+    "_attention_prefill_ragged",
+    "_attention_prefill_ragged_cpu",
+    "_attention_sequence_prefill",
+    "_attention_sequence_prefill_with_mask",
+    "_compact_kv_copy",
+    "_compact_kv_copy_cpu",
+    "_copy_single_page",
+    "_copy_single_page_cpu",
+    "_copy_single_page_mla",
+    "_kv_cache_debug_get_kv",
+    "_kv_cache_debug_get_kv_mla",
+    "_kv_cache_transpose_append",
+    "_kv_cache_transpose_append_mla",
+    "_merge_state_inplace",
+    "_merge_state_inplace_cpu",
+    "llama_rope_with_position_map",
+    "tree_attn",
+    "tree_attn_cpu",
+    "tree_attn_with_paged_kv_cache",
+    "tree_attn_with_paged_kv_cache_cpu",
+]
+
+
+class PagedKVCache(Object):  # pylint: disable=too-few-public-methods
+    """The Paged KV Cache used in LLM batching for efficient attention computation."""
+
+    extern_mods: list[tvm.runtime.Module] = []
+
+    def attention_with_fused_qkv(
+        self,
+        layer_id: int,
+        qkv: Tensor,
+        num_qo_heads: int,
+        sm_scale: float,
+    ) -> Tensor:
+        """Compute attention with the given fused q/k/v data and in-cache k/v data
+        on the specified layer. Rotary position embeddings are applied to k/v
+        within this function.
+
+        - For prefill, the input qkv and output tensor have shape
+        (1, total_seq_len) for the first two dimensions.
+        - For decode, the input qkv and output tensor have shape
+        (batch_size, 1) for the first two dimensions.
+        - The input qkv have `2 * num_qo_heads + num_kv_heads` at the third dim.
+        - The output tensor have `num_qo_heads` at the third dim.
+        - The input qkv and output tensor have `head_dim` at the last dim.
+        """
+        # pylint: disable=protected-access
+        b, s, _, d = qkv._expr.ty.shape
+        qkv = qkv.reshape(b * s, qkv.shape[2], d)
+        return Tensor(
+            _expr=rx.BlockBuilder.current().emit(
+                rx.call_dps_packed(
+                    "vm.builtin.attention_kv_cache_attention_with_fused_qkv",
+                    [
+                        self._expr,
+                        rx.prim_value(layer_id),  # type: ignore[arg-type]
+                        rx.prim_value(sm_scale),
+                        qkv._expr,
+                    ],
+                    out_ty=rx.TensorType((b * s, num_qo_heads, d), qkv.dtype),
+                )
+            )
+        ).reshape(b, s, num_qo_heads, d)
+
+    def self_attention(  # pylint: disable=too-many-locals
+        self,
+        layer_id: int,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        sm_scale: float,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that computes ragged self attention with Q/K/V data."""
+        # pylint: disable=protected-access
+        b, s, h_qo, d_qk = q._expr.ty.shape
+        _, _, h_kv, d_v = v._expr.ty.shape
+        q = q.reshape(b * s, h_qo, d_qk)
+        k = k.reshape(b * s, h_kv, d_qk)
+        v = v.reshape(b * s, h_kv, d_v)
+        bb = rx.BlockBuilder.current()
+        attn_results = bb.emit(
+            rx.call_dps_packed(
+                "vm.builtin.attention_kv_cache_self_attention",
+                [
+                    self._expr,
+                    rx.prim_value(layer_id),  # type: ignore[arg-type]
+                    rx.prim_value(sm_scale),
+                    q._expr,
+                    k._expr,
+                    v._expr,
+                ],
+                out_ty=[
+                    rx.TensorType((b * s, h_qo, d_v), q.dtype),
+                    rx.TensorType((b * s, h_qo), "float32"),
+                ],
+            )
+        )
+        assert isinstance(attn_results.ty, rx.TupleType)
+        assert len(attn_results.ty.fields) == 2
+        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, d_v)
+        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
+        return o, lse
+
+    def cross_attention(
+        self,
+        layer_id: int,
+        q: Tensor,
+        v_head_dim: int,
+        sm_scale: float,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that computes paged cross attention with Q and in-cache KV data."""
+        # pylint: disable=protected-access
+        b, s, h_qo, d_qk = q._expr.ty.shape
+        q = q.reshape(b * s, h_qo, d_qk)
+        bb = rx.BlockBuilder.current()
+        attn_results = bb.emit(
+            rx.call_dps_packed(
+                "vm.builtin.attention_kv_cache_cross_attention",
+                [
+                    self._expr,
+                    rx.prim_value(layer_id),  # type: ignore[arg-type]
+                    rx.prim_value(sm_scale),
+                    q._expr,
+                ],
+                out_ty=[
+                    rx.TensorType((b * s, h_qo, v_head_dim), q.dtype),
+                    rx.TensorType((b * s, h_qo), "float32"),
+                ],
+            )
+        )
+        assert isinstance(attn_results.ty, rx.TupleType)
+        assert len(attn_results.ty.fields) == 2
+        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, v_head_dim)
+        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
+        return o, lse
+
+    def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
+        """Fine-grained API that appends the MLA K/V data to KV cache."""
+        # pylint: disable=protected-access
+        b, s, _, d_qk = kv._expr.ty.shape
+        kv = kv.reshape(b * s, d_qk)
+        return PagedKVCache(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.attention_kv_cache_append_mla_kv",
+                self._expr,
+                rx.prim_value(layer_id),  # type: ignore[arg-type]
+                kv._expr,
+                ty_args=rx.AnyType(),
+            ),
+            _name="paged_kv_cache",
+        )
+
+    def merge_attn_output_inplace(
+        self,
+        o_self_attn: Tensor,
+        lse_self_attn: Tensor,
+        o_cross_attn: Tensor,
+        lse_cross_attn: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that merges the attention output from two sources.
+        The first two tensors will be inplace updated.
+        """
+        # pylint: disable=protected-access
+        b, s, h_qo, d_v = o_self_attn._expr.ty.shape
+        o_self_attn = o_self_attn.reshape(b * s, h_qo, d_v)
+        lse_self_attn = lse_self_attn.reshape(b * s, h_qo)
+        o_cross_attn = o_cross_attn.reshape(b * s, h_qo, d_v)
+        lse_cross_attn = lse_cross_attn.reshape(b * s, h_qo)
+        bb = rx.BlockBuilder.current()
+        merge_results = bb.emit(
+            rx.call_pure_packed(
+                "vm.builtin.attention_kv_cache_merge_attn_output_inplace",
+                self._expr,
+                o_self_attn._expr,
+                lse_self_attn._expr,
+                o_cross_attn._expr,
+                lse_cross_attn._expr,
+                ty_args=rx.TupleType(
+                    [o_self_attn._expr.ty, lse_self_attn._expr.ty]
+                ),
+            )
+        )
+        assert isinstance(merge_results.ty, rx.TupleType)
+        assert len(merge_results.ty.fields) == 2
+        o_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 0))).reshape(
+            b, s, h_qo, d_v
+        )
+        lse_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 1))).reshape(b, s, h_qo)
+        return o_self_attn, lse_self_attn
+
+    def get_query_positions(self, total_length: tirx.Expr) -> Tensor:
+        """Get the in-sequence positions of each slot in the query,
+        which are needed for applying positional embeddings in some models.
+
+        Parameters
+        ----------
+        total_length : tirx.Expr
+            The summed-up total sequence length of queries in
+            the batch being forwarded.
+
+        Returns
+        -------
+        q_positions : Tensor
+            The in-sequence query positions, in shape `(total_length,)`
+        """
+        return Tensor(
+            _expr=rx.BlockBuilder.current().emit(
+                rx.call_pure_packed(
+                    "vm.builtin.attention_kv_cache_get_query_positions",
+                    self._expr,
+                    ty_args=rx.TensorType((total_length,), "int32"),
+                )
+            )
+        )
+
+    # pylint: enable=protected-access
+
+
+def _prepare_yarn_rope_scaling(rope_scaling: dict[str, Any] | None, rope_theta: float | None) -> dict[str, Any] | None:
+    """Ensure Yarn-specific scaling configs include the theta metadata."""
+    if rope_scaling is None:
+        return None
+    if rope_scaling.get("rope_type") != "yarn":
+        return rope_scaling
+
+    rope_scaling_updated = dict(rope_scaling)
+    if "inv_theta_log_scale" not in rope_scaling_updated and rope_theta is not None:
+        theta_value = float(rope_theta)
+        rope_scaling_updated["inv_theta_log_scale"] = 1.0 / (2 * math.log(theta_value))
+    return rope_scaling_updated
+
+
+class FlashInferPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
+    """Paged KV cache using FlashInfer (CUDA) kernels."""
+
+    def __init__(  # pylint: disable=too-many-locals
+        self,
+        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
+        max_batch_size: tirx.Var,
+        max_total_seq_len: tirx.Var,
+        prefill_chunk_size: tirx.Var,
+        page_size: tirx.Var,
+        support_sliding_window: tirx.Var,
+        layer_partition: rx.ShapeExpr,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        mla_original_qk_head_dim: int,
+        mla_original_v_head_dim: int,
+        rope_mode: RopeMode,
+        rope_scale: int,
+        rope_theta: int,
+        rope_scaling: dict[str, Any],
+        rope_ext_factors: rx.Expr,
+        rotary_dim: int,
+        enable_disaggregation: bool,
+        dtype: str,
+        target: Target,
+        name: str = "paged_kv_cache",
+    ) -> None:
+        """Create a paged KV cache object with FlashInfer kernels.
+
+        Parameters
+        ----------
+        max_batch_size : tirx.Var
+            The maximum allowed batch size of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        max_total_seq_len : tirx.Var
+            The maximum allowed total sequence length of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        prefill_chunk_size : tirx.Var
+            The maximum total sequence length in a prefill.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        page_size : tirx.Var
+            The size (a.k.a. number of tokens) of each page.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        support_sliding_window : tirx.Var
+            0 or 1, denoting whether the KV cache supports sliding window.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        layer_partition : rx.ShapeExpr
+            The KV cache layer partition for pipeline stages.
+            It is an indptr array, denoting the starting layer of each pipeline stage.
+        rope_mode : RopeMode
+            The RoPE mode of the Paged KV cache.
+            If it is normal, RoPE will be applied to k before adding k to cache.
+            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+        rope_scale : int
+            The scale of rotary position embedding.
+        rope_theta : int
+            The base of rotary position embedding.
+        rope_scaling: Dict[str, Any]
+            The RoPE scaling information dict.
+        rope_ext_factors: rx.Expr
+            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
+        rotary_dim : int
+            The number of dimensions in the embedding that RoPE is applied to.
+        enable_disaggregation : bool
+            Whether to enable disaggregation in the KV cache.
+        """
+        assert rope_mode != RopeMode.INLINE, "FlashInfer RoPE does not support inline mode."
+        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
+
+        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
+        if attn_kind_single == "mha_sliding":
+            attn_kind_single = "mha"
+        flashinfer_prefill_mods = rx.backend.cuda.flashinfer.gen_flashinfer_prefill_module(
+            dtype_q=dtype,
+            dtype_kv=dtype,
+            dtype_o=dtype,
+            qk_head_dim=(qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim),
+            v_head_dim=(v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim),
+            enable_inline_rope=False,
+            return_static_libs=True,
+        )
+        flashinfer_decode_mods = (
+            rx.backend.cuda.flashinfer.gen_flashinfer_decode_module(
+                dtype_q=dtype,
+                dtype_kv=dtype,
+                dtype_o=dtype,
+                qk_head_dim=qk_head_dim,
+                v_head_dim=v_head_dim,
+                enable_inline_rope=False,
+                return_static_libs=True,
+            )
+            if attn_kind_single == "mha"
+            else []
+        )
+        flashinfer_mla_mods = (
+            rx.backend.cuda.flashinfer.gen_flashinfer_mla_module(
+                dtype_q=dtype,
+                dtype_kv=dtype,
+                dtype_o=dtype,
+                head_dim_ckv=v_head_dim,
+                head_dim_kpe=qk_head_dim - v_head_dim,
+                return_static_libs=True,
+            )
+            if attn_kind_single == "mla"
+            else []
+        )
+        self.extern_mods = flashinfer_prefill_mods + flashinfer_decode_mods + flashinfer_mla_mods
+
+        bb = rx.BlockBuilder.current()
+        mha_functions = (
+            [
+                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
+                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+            ]
+            if attn_kind_single == "mha"
+            else [rx.Tuple([]) for _ in range(6)]
+        )
+        ragged_prefill_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
+        mla_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
+        attn_merge_functions = [
+            bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
+        ]
+        if attn_kind_single == "mla":
+            attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
+
+        if isinstance(attn_kind, list):
+            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
+        else:
+            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
+
+        args = [
+            rx.ShapeExpr(
+                [
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ]
+            ),
+            layer_partition,
+            rx.prim_value(num_attention_heads),
+            rx.prim_value(num_key_value_heads),
+            rx.prim_value(qk_head_dim),
+            rx.prim_value(v_head_dim),
+            rx.ShapeExpr(attn_kind),
+            rx.prim_value(enable_disaggregation),
+            rx.prim_value(rope_mode),
+            rx.prim_value(rope_scale),
+            rx.prim_value(rope_theta),
+            rope_ext_factors,
+            rx.op.zeros((), dtype),
+            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
+            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
+            ragged_prefill_function,
+            *mha_functions,
+            mla_function,
+            rx.Tuple(attn_merge_functions),
+            bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+            bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
+            bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+            bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
+        ]
+        super().__init__(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.paged_attention_kv_cache_create",
+                *args,
+                ty_args=rx.AnyType(),
+            ),
+            _name=name,
+        )
+
+
+class TIRPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
+    """Paged KV cache using TIR kernels."""
+
+    def __init__(  # pylint: disable=too-many-locals
+        self,
+        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
+        max_batch_size: tirx.Var,
+        max_total_seq_len: tirx.Var,
+        prefill_chunk_size: tirx.Var,
+        page_size: tirx.Var,
+        support_sliding_window: tirx.Var,
+        layer_partition: rx.ShapeExpr,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        mla_original_qk_head_dim: int,
+        mla_original_v_head_dim: int,
+        rope_mode: RopeMode,
+        rope_scale: int,
+        rope_theta: int,
+        rope_scaling: dict[str, Any],
+        rope_ext_factors: rx.Expr,
+        rotary_dim: int,
+        enable_disaggregation: bool,
+        dtype: str,
+        target: Target,
+        name: str = "paged_kv_cache",
+    ) -> None:
+        """Create a paged KV cache object with TIR kernels.
+
+        Parameters
+        ----------
+        max_batch_size : tirx.Var
+            The maximum allowed batch size of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        max_total_seq_len : tirx.Var
+            The maximum allowed total sequence length of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        prefill_chunk_size : tirx.Var
+            The maximum total sequence length in a prefill.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        page_size : tirx.Var
+            The size (a.k.a. number of tokens) of each page.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        support_sliding_window : tirx.Var
+            0 or 1, denoting whether the KV cache supports sliding window.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        layer_partition : rx.ShapeExpr
+            The KV cache layer partition for pipeline stages.
+            It is an indptr array, denoting the starting layer of each pipeline stage.
+        rope_mode : RopeMode
+            The RoPE mode of the Paged KV cache.
+            If it is normal, RoPE will be applied to k before adding k to cache.
+            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+        rope_scale : int
+            The scale of rotary position embedding.
+        rope_theta : int
+            The base of rotary position embedding.
+        rope_scaling: Dict[str, Any]
+            The RoPE scaling information dict.
+        rope_ext_factors: rx.Expr
+            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
+        rotary_dim : int
+            The number of dimensions in the embedding that RoPE is applied to.
+        enable_disaggregation : bool
+            Whether to enable disaggregation in the KV cache.
+        target : Target
+            The target to build the model to.
+        """
+        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
+        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
+        if attn_kind_single == "mha_sliding":
+            attn_kind_single = "mha"
+        if isinstance(attn_kind, list):
+            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
+        else:
+            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
+        bb = rx.BlockBuilder.current()
+        args = [
+            rx.ShapeExpr(
+                [
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ]
+            ),
+            layer_partition,
+            rx.prim_value(num_attention_heads),
+            rx.prim_value(num_key_value_heads),
+            rx.prim_value(qk_head_dim),
+            rx.prim_value(v_head_dim),
+            rx.ShapeExpr(attn_kind),
+            rx.prim_value(enable_disaggregation),
+            rx.prim_value(rope_mode),
+            rx.prim_value(rope_scale),
+            rx.prim_value(rope_theta),
+            rope_ext_factors,
+            rx.op.zeros((), dtype),
+            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
+            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
+        ]
+
+        if target.kind.name == "llvm":
+            if attn_kind_single == "mla":
+                raise ValueError("MLA is not supported in TIR kernels for now.")
+            args.extend(
+                [
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_prefill_cpu_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
+                    rx.Tuple([]),  # f_mla_prefill
+                    rx.Tuple([bb.add_func(_merge_state_inplace_cpu(dtype), "tir_attention_merge_state_cpu")]),
+                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+                    bb.add_func(_copy_single_page_cpu(num_key_value_heads, page_size, qk_head_dim, dtype), "kv_cache_copy_single_page_cpu"),
+                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+                    bb.add_func(_compact_kv_copy_cpu(num_key_value_heads, qk_head_dim, dtype), "kv_cache_compact_kv_copy_cpu"),
+                ]
+            )
+        else:
+            ragged_qk_head_dim = qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim
+            ragged_v_head_dim = v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim
+            args.append(rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
+            mha_functions = (
+                [
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+                ]
+                if attn_kind_single == "mha"
+                else [rx.Tuple([]) for _ in range(6)]
+            )
+            mla_function = rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
+            attn_merge_functions = [
+                bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
+            ]
+            if attn_kind_single == "mla":
+                attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
+            args.extend(mha_functions)
+            args.append(mla_function)
+            args.extend(
+                [
+                    rx.Tuple(attn_merge_functions),
+                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+                    bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
+                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+                    bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
+                ]
+            )
+
+        super().__init__(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.paged_attention_kv_cache_create",
+                *args,
+                ty_args=rx.AnyType(),
+            ),
+            _name=name,
+        )

python/tvm/relax/frontend/nn/llm/position_embedding.py

@@ -0,0 +1,894 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Operators for positional embeddings, e.g. RoPE."""
+
+import math
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+
+from tvm import tirx
+from tvm.relax.frontend.nn import Tensor, op
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name
+
+
+def rope_freq_default(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE and then return the cosine and sine of it.
+
+    Parameters
+    ----------
+    s : tirx.Var
+        The position index.
+
+    d : tirx.Var
+        The dimension index.
+
+    d_range : int
+        The maximum dimension index.
+
+    theta : float
+        The theta value in RoPE, which controls the frequency.
+
+    dtype : str
+        The data type of the output.
+
+    Returns
+    -------
+    cos_freq : Tensor
+        The cosine of the inverse frequency.
+
+    sin_freq : Tensor
+        The sine of the inverse frequency.
+
+    var_map: Dict[tirx.Var, tirx.Expr]
+        The common expression map.
+    """
+    freq = s / tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_gptj(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE for gptj RoPE scaling."""
+    freq = s / tirx.power(theta, 2 * (d // 2) % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_llama4(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama4 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, 2 * (d // 2) / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+
+    llama4_inv_scaling_factor = 1.0 / factor
+
+    if high_freq_factor == low_freq_factor:
+        wavelength = tirx.const(2 * math.pi, "float32") / orig_freq_var
+        threshold_wavelen = tirx.const(
+            original_max_position_embeddings / low_freq_factor, "float32"
+        )
+
+        scaled_freq = tirx.if_then_else(
+            wavelength > threshold_wavelen, orig_freq_var / factor, orig_freq_var
+        )
+        smoothed_freq = s * scaled_freq
+
+    else:
+        # Original smooth interpolation logic
+        inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+
+        llama4_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+        llama4_beta = low_freq_factor * inv_diff_freq_factor
+        smooth = tirx.max(0.0, tirx.min(1.0, llama4_alpha * orig_freq_var - llama4_beta))
+        smoothed_freq = s * (
+            (1.0 - smooth) * orig_freq_var * llama4_inv_scaling_factor + smooth * orig_freq_var
+        )
+
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_llama3(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama3 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, d * 2 % d_range / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+    inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+    llama3_inv_scaling_factor = 1.0 / factor
+    llama3_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+    llama3_beta = low_freq_factor * inv_diff_freq_factor
+    smooth = tirx.max(0.0, tirx.min(1.0, llama3_alpha * orig_freq_var - llama3_beta))
+    smoothed_freq = s * (
+        (1.0 - smooth) * orig_freq_var * llama3_inv_scaling_factor + smooth * orig_freq_var
+    )
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_longrope(  # pylint: disable=too-many-arguments
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    max_position_embeddings: int,
+    original_max_position_embeddings: int,
+    ext_factors: T.Buffer | None = None,
+):
+    """Compute the inverse frequency of RoPE for longrope scaling."""
+    scale = max_position_embeddings / original_max_position_embeddings
+    scaling_factor = (
+        math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings))
+        if scale > 1.0
+        else 1.0
+    )
+    divisor = tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    if ext_factors is not None:
+        divisor = ext_factors[d % (d_range // 2)] * divisor
+    freq = s / divisor
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = (tirx.cos(freq_var) * scaling_factor).astype(dtype)
+    sin_freq = (tirx.sin(freq_var) * scaling_factor).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def yarn_find_correction_dim(
+    num_rotations: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Inverse dim formula to find dim based on number of rotations"""
+    return (
+        d * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) * inv_theta_log_scale
+    )
+
+
+def yarn_find_correction_range(
+    low_rot: int,
+    high_rot: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Find the correction range based on the number of rotations"""
+    low = yarn_find_correction_dim(
+        low_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    high = yarn_find_correction_dim(
+        high_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    return tirx.max(low, 0), tirx.min(high, d - 1)
+
+
+def rope_freq_yarn(
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float | tirx.Expr,
+    dtype: str,
+    original_max_position_embeddings: int,
+    scaling_factor: float,
+    beta_fast: int,
+    beta_slow: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):  # pylint: disable=too-many-arguments, too-many-locals
+    """Compute the inverse frequency of RoPE for yarn RoPE scaling."""
+
+    exponent = d * 2 % d_range / tirx.const(d_range, "float32")
+    freq_power = tirx.power(theta, exponent)
+    freq_extra = tirx.const(1, "float32") / freq_power
+    freq_inter = tirx.const(1, "float32") / (scaling_factor * freq_power)
+
+    low, high = yarn_find_correction_range(
+        beta_fast,
+        beta_slow,
+        d_range,
+        original_max_position_embeddings,
+        inv_theta_log_scale=inv_theta_log_scale,
+    )
+    high = tirx.if_then_else(low == high, high + 0.001, high)
+    inv_freq_mask = tirx.const(1, "float32") - tirx.max(
+        tirx.min((d - low) / (high - low), 1.0), 0.0
+    ).astype("float32")
+    inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask
+    freq = s * inv_freq
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def switch_rope_freq_func(rope_scaling: dict[str, Any]) -> Callable:
+    """Return the RoPE inverse frequency computation function based
+    on the given RoPE scaling.
+    """
+    if "rope_type" not in rope_scaling:
+        return rope_freq_default
+    if rope_scaling["rope_type"] == "gptj":
+        return rope_freq_gptj
+    if rope_scaling["rope_type"] == "llama3":
+        return partial(
+            rope_freq_llama3,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "llama4":
+        return partial(
+            rope_freq_llama4,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "longrope":
+        return partial(
+            rope_freq_longrope,
+            max_position_embeddings=rope_scaling["max_position_embeddings"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "yarn":
+        inv_theta_log_scale = rope_scaling.get("inv_theta_log_scale")
+        assert inv_theta_log_scale is not None, "inv_theta_log_scale must be precomputed for YaRN"
+        return partial(
+            rope_freq_yarn,
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+            scaling_factor=rope_scaling["factor"],
+            beta_fast=rope_scaling["beta_fast"],
+            beta_slow=rope_scaling["beta_slow"],
+            inv_theta_log_scale=inv_theta_log_scale,
+        )
+    raise ValueError(f"Unsupported RoPE scaling type: {rope_scaling['rope_type']}")
+
+
+# mypy: disable-error-code="attr-defined"
+
+
+def llama_rope(  # pylint: disable=too-many-arguments
+    qkv: Tensor,
+    total_seq_len: tirx.Var,
+    theta: float,
+    scale: float,
+    num_q_heads: int,
+    num_kv_heads: int,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Llama-style RoPE. Given a fused QKV tensor, it returns three tensors, Q, K, and V, where Q
+    and K are rotated by RoPE while V remains unchanged.
+
+    Parameters
+    ----------
+    qkv : Tensor
+        The fused QKV tensor of shape: [batch_size, seq_len, #q_heads + #kv_heads * 2, head_dim]
+
+    total_seq_len : tirx.Var
+        The total sequence length after being concatenated with KVCache. It is used to compute the
+        offset of RoPE.
+
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : Optional[int]
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+
+    Returns
+    -------
+    q : Tensor
+        The query tensor of shape [batch_size, seq_len, #q_heads, head_dim] w/ RoPE applied
+
+    k : Tensor
+        The key tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/ RoPE applied
+
+    v : Tensor
+        The value tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/o RoPE applied
+    """
+    _, _, fused_heads, head_dim = qkv.shape
+    assert fused_heads == num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    dtype = qkv.dtype
+    scale = tirx.const(scale, dtype)
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        b: tirx.Var,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        offset: tirx.Var,
+    ):
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            (s + offset) * scale, d, rotary_dim, theta, dtype
+        )
+        cos = cos_freq * x[b, s, h, d]
+        if rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[b, s, h, d + 1],
+                x[b, s, h, d - 1],
+            )
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[b, s, h, d + rotary_dim // 2],
+                x[b, s, h, d - rotary_dim // 2],
+            )
+        expr = cos + sin
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        total_seq_len: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        batch_size = T.int64()
+        seq_len = T.int64()
+        qkv = T.match_buffer(var_qkv, (batch_size, seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (batch_size, seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        for iters in T.grid(batch_size, seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                b, s, h, d = T.axis.remap("SSSS", iters)
+                if h < num_q_heads:
+                    q[b, s, h, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[b, s, h - num_q_heads, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                else:
+                    v[b, s, h - (num_q_heads + num_kv_heads), d] = qkv[b, s, h, d]
+
+    b, s, _, _ = qkv.shape
+    return op.tensor_ir_op(  # pylint: disable=no-member
+        fused_rope,
+        "llama_rope",
+        args=[qkv, total_seq_len],
+        out=(
+            Tensor.placeholder((b, s, num_q_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+        ),
+    )
+
+
+def llama_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[s, h, d + rotary_dim // 2],
+                x[s, h, d - rotary_dim // 2],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope
+
+
+def llama4_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            # Data layout is different for llama4 vs llama3
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope

python/tvm/relax/frontend/nn/spec.py

@@ -0,0 +1,257 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Compilation specifications, for example, dynamic shape inputs."""
+
+import inspect
+import typing
+
+if typing.TYPE_CHECKING:
+    from .core import Module as nn_module_class
+
+ArgSpecType = typing.Union["Int", "Tensor"]
+MethodSpecType = typing.Union["MethodSpec", dict[str, ArgSpecType]]
+ModuleSpecType = typing.Union["ModuleSpec", dict[str, MethodSpecType]]
+SpecAny = typing.Union["Object", "Int", "Tensor", "Tuple"]
+
+
+class Int:  # pylint: disable=too-few-public-methods
+    """An integer input"""
+
+    def __init__(self) -> None:
+        pass
+
+    def __repr__(self) -> str:
+        return "int"
+
+
+class Tensor:  # pylint: disable=too-few-public-methods
+    """A tensor input with static ndim and dtype, but can have symbolic shapes."""
+
+    shape: list[int | str]
+    dtype: str
+
+    def __init__(self, shape: typing.Sequence[int | str], dtype: str) -> None:
+        self.shape = list(shape)
+        self.dtype = dtype
+
+    def __repr__(self) -> str:
+        shape = ", ".join(str(i) for i in self.shape)
+        return f"Tensor([{shape}], '{self.dtype}')"
+
+
+class Object:  # pylint: disable=too-few-public-methods
+    """An non-tensor opaque frontend object."""
+
+    object_type: type
+
+    def __init__(self, object_type: type) -> None:
+        self.object_type = object_type
+
+    def __repr__(self) -> str:
+        return "object"
+
+
+class Tuple:  # pylint: disable=too-few-public-methods
+    """A tuple input or a list input"""
+
+    name: str
+    elements: list[SpecAny] | tuple[SpecAny, ...]
+
+    def __init__(
+        self,
+        name: str,
+        elements: list[SpecAny] | tuple[SpecAny, ...],
+    ) -> None:
+        assert isinstance(elements, tuple | list), f"Unsupported container type: {type(elements)}"
+        self.name = name
+        self.elements = elements
+
+    def __repr__(self) -> str:
+        return self.elements.__repr__()
+
+
+class MethodSpec:
+    """A spec for a compiled method"""
+
+    method: typing.Callable
+    arg_names: list[str]
+    arg_specs: list[ArgSpecType]
+    param_mode: str  # "plain", "packed", "none"
+    effect_mode: str  # "plain", "packed", "none"
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        method: typing.Callable,
+        arg_names: list[str],
+        arg_specs: list[ArgSpecType],
+        param_mode: str,
+        effect_mode: str,
+    ):
+        if param_mode not in ["plain", "packed", "none"]:
+            raise ValueError(f"Invalid param_mode: {param_mode}")
+        if effect_mode not in ["plain", "packed", "none"]:
+            raise ValueError(f"Invalid effect_mode: {effect_mode}")
+        self.method = method
+        self.arg_names = arg_names
+        self.arg_specs = arg_specs
+        self.param_mode = param_mode
+        self.effect_mode = effect_mode
+
+    def _repr(self, name: str) -> str:
+        args = ", ".join(
+            f"{name}: {spec}"
+            for name, spec in zip(
+                self.arg_names,
+                self.arg_specs,
+            )
+        )
+        return f"{name}({args})"
+
+    def __repr__(self) -> str:
+        return self._repr(name="MethodSpec")
+
+    @staticmethod
+    def from_raw(spec: MethodSpecType, method: typing.Callable) -> "MethodSpec":
+        """Create MethodSpec from raw python dictionaries.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            MethodSpec.from_raw(
+                spec={
+                    "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
+                    "total_seq_len": "int",
+                },
+                method=module.prefill,
+            )
+        """
+        if isinstance(spec, MethodSpec):
+            return spec
+        config: dict[str, typing.Any] = spec.pop("$", {})  # type: ignore[assignment]
+        param_mode = config.get("param_mode", "plain")
+        effect_mode = config.get("effect_mode", "plain")
+        method_signature = inspect.signature(method)
+        arg_names = list(method_signature.parameters.keys())
+        arg_specs = []
+
+        def _convert_arg_spec(arg_spec, arg_name):
+            if arg_spec is Int or arg_spec is int:
+                return Int()
+            if isinstance(arg_spec, str) and arg_spec == "int":
+                return Int()
+            if isinstance(arg_spec, Int | Tensor | Object):
+                return arg_spec
+            if isinstance(arg_spec, tuple | list | Tuple):
+                return Tuple(
+                    arg_name,
+                    elements=type(arg_spec)(
+                        [
+                            _convert_arg_spec(arg_spec_i, f"{arg_name}_{i}")
+                            for i, arg_spec_i in enumerate(arg_spec)
+                        ]
+                    ),
+                )
+            raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")
+
+        for arg_name in arg_names:
+            if arg_name in spec:
+                arg_spec = spec[arg_name]
+                arg_spec = _convert_arg_spec(arg_spec, arg_name)
+                arg_specs.append(arg_spec)
+        return MethodSpec(
+            method,
+            arg_names,
+            arg_specs,
+            param_mode=param_mode,
+            effect_mode=effect_mode,
+        )
+
+    @staticmethod
+    def from_torch(args: list[typing.Any], method: typing.Callable) -> "MethodSpec":
+        """Converts a list of torch tensors to MethodSpec."""
+        from .torch import (  # pylint: disable=import-outside-toplevel
+            _method_spec_from_torch,
+        )
+
+        return _method_spec_from_torch(args, method)
+
+
+class ModuleSpec:
+    """A spec for a compiled nn.Module"""
+
+    module: "nn_module_class"
+    method_names: list[str]
+    method_specs: list[MethodSpec]
+
+    def __init__(
+        self,
+        module: "nn_module_class",
+        method_names: list[str],
+        method_specs: list[MethodSpec],
+    ) -> None:
+        self.module = module
+        self.method_names = method_names
+        self.method_specs = method_specs
+
+    @staticmethod
+    def from_raw(spec: ModuleSpecType, module: "nn_module_class") -> "ModuleSpec":
+        """Create ModuleSpec from raw python dictionaries.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            ModuleSpec.from_raw(
+                spec={
+                    "prefill": {
+                        "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
+                        "total_seq_len": int,
+                    },
+                    "decode": {
+                        "inputs": spec.Tensor([batch_size, 1], "int32"),
+                        "total_seq_len": int,
+                    },
+                    "softmax_with_temperature": {
+                        "logits": spec.Tensor([1, 1, config.vocab_size], "float32"),
+                        "temperature": spec.Tensor([], "float32"),
+                    },
+                },
+                module=module,
+            )
+        """
+        if isinstance(spec, ModuleSpec):
+            return spec
+        method_names = list(spec.keys())
+        method_specs: list[MethodSpec] = []
+        for method_name in method_names:
+            method_spec = spec[method_name]
+            if isinstance(method_spec, MethodSpec):
+                pass
+            else:
+                method_spec = MethodSpec.from_raw(method_spec, getattr(module, method_name))
+            method_specs.append(method_spec)
+        return ModuleSpec(module, method_names, method_specs)
+
+    def __repr__(self) -> str:
+        return "ModuleSpec:\n" + "\n".join(
+            "  " + spec._repr(name)  # pylint: disable=protected-access
+            for name, spec in zip(
+                self.method_names,
+                self.method_specs,
+            )
+        )

python/tvm/relax/frontend/nn/subroutine.py

@@ -0,0 +1,188 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=too-many-lines,invalid-name,protected-access
+"""nn.Module mixin for subroutine dispatch"""
+
+import collections
+import contextlib
+import functools
+import inspect
+import re
+import typing
+
+import tvm_ffi
+
+from tvm import ir, relax
+from tvm.relax.frontend import nn
+
+
+def _camel_to_snake(name):
+    """Convert from CamelCase to snake_case"""
+
+    # Adapted from https://stackoverflow.com/a/1176023
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", name)
+    name = name.lower()
+    return name
+
+
+def _normalize_expr(block_builder, arg, as_relax_expr=False):
+    """Ensure that an argument is a relax.Expr with type"""
+    if isinstance(arg, tuple):
+        arg = relax.Tuple([_normalize_expr(block_builder, element) for element in arg])
+
+    if isinstance(arg, relax.Expr) and arg.ty.is_missing():
+        arg = block_builder.emit(arg)
+
+    if isinstance(arg, nn.Tensor) and as_relax_expr:
+        arg = arg._expr
+
+    return arg
+
+
+def _get_ty(arg):
+    if isinstance(arg, relax.Expr):
+        return arg.ty
+    elif isinstance(arg, nn.Tensor):
+        return arg._expr.ty
+    elif isinstance(arg, tuple | list | tvm_ffi.Array):
+        return relax.TupleType([_get_ty(field) for field in arg])
+    else:
+        raise TypeError(f"Cannot find type for {arg} of type {type(arg)}")
+
+
+class SubroutineMixin:
+    """A mixin that generates a
+
+    Contains common logic for `tvm.relax.frontend.nn.Module` and
+    `tvm.relax.testing.nn.Module`.
+    """
+
+    define_subroutine: bool = False
+
+    def __init_subclass__(cls):
+        """Update the cls.forward of subclasses"""
+        if hasattr(cls, "forward"):
+            is_wrapped = getattr(cls.forward, "_is_subroutine_mixin", False)
+            if not is_wrapped:
+                cls.forward = cls._subroutine_dispatch(cls.forward)
+
+    @classmethod
+    def _subroutine_dispatch(cls, old_forward):
+        @functools.wraps(old_forward)
+        def new_forward(self, *args, **kwargs):
+            if not self.define_subroutine:
+                return old_forward(self, *args, **kwargs)
+
+            block_builder = relax.BlockBuilder.current()
+            assert block_builder is not None, (
+                f"Class {type(self)} has cls.define_subroutines = True, "
+                "but is called outsdie of a block_builder environment.  "
+                "relax.BlockBuilder.current() is required "
+                "to determine where to generate the subroutine."
+            )
+
+            func_args = self._normalize_subroutine_args(block_builder, *args, **kwargs)
+            subroutine, is_nn_tensor_output = self._get_subroutine(
+                block_builder, old_forward, func_args
+            )
+            subroutine_args = [
+                arg._expr if isinstance(arg, nn.Tensor) else arg
+                for arg in [*func_args.values(), *self.parameters()]
+            ]
+
+            out = subroutine(*subroutine_args)
+
+            if is_nn_tensor_output:
+                if out.ty.is_missing():
+                    out = block_builder.emit(out, name_hint=f"{subroutine.name_hint}_output")
+                out = nn.Tensor(_expr=out)
+            return out
+
+        new_forward._is_subroutine_mixin = True
+
+        return new_forward
+
+    def _normalize_subroutine_args(
+        self, block_builder, *args, **kwargs
+    ) -> typing.OrderedDict[str, relax.Expr]:
+        signature = inspect.signature(self.forward)
+        bindings = signature.bind(*args, **kwargs)
+        func_args = collections.OrderedDict(
+            (name, _normalize_expr(block_builder, arg)) for name, arg in bindings.arguments.items()
+        )
+        return func_args
+
+    def _get_subroutine(
+        self,
+        block_builder,
+        old_forward: typing.Callable,
+        func_args: typing.OrderedDict[str, relax.Expr],
+    ) -> (ir.GlobalVar, bool):
+        cls = type(self)
+        if not hasattr(cls, "_gvar"):
+            cls._gvar = {}
+
+        model_params = [
+            param._expr if isinstance(param, nn.Tensor) else param for param in self.parameters()
+        ]
+
+        arg_ty = _get_ty([*func_args.values(), *model_params])
+        is_dataflow = block_builder.current_block_is_dataflow()
+        lookup_key = (
+            old_forward,
+            tvm_ffi.structural_hash(arg_ty, map_free_vars=True),
+            is_dataflow,
+        )
+
+        for cached_ty, cached_result in cls._gvar.get(lookup_key, []):
+            if tvm_ffi.structural_equal(cached_ty, arg_ty, map_free_vars=True):
+                return cached_result
+
+        func_name = _camel_to_snake(cls.__name__)
+        func_params = [relax.Var(name, ty) for name, ty in zip(func_args, arg_ty.fields)]
+        old_forward_args = [
+            nn.Tensor(_expr=param) if isinstance(old_arg, nn.Tensor) else param
+            for param, old_arg in zip(func_params, func_args.values())
+        ]
+
+        with block_builder.function(func_name, [*func_params, *model_params], private=True):
+            with contextlib.ExitStack() as stack:
+                if is_dataflow:
+                    stack.enter_context(block_builder.dataflow())
+
+                out = old_forward(self, *old_forward_args)
+                is_nn_tensor_output = isinstance(out, nn.Tensor)
+                if is_nn_tensor_output:
+                    out = out._expr
+
+                if is_dataflow:
+                    out = block_builder.emit_output(out)
+            gvar = block_builder.emit_func_output(out)
+
+        # The relax.Var instances in model_params, along with any
+        # tirx.Var instances in the type, appear in both the
+        # calling scope and as parameters for the subroutine.  To
+        # maintain SSA, replace all relax and TIR variables in the
+        # subroutine.
+        mod = block_builder.get()
+        mod.update_func(gvar, relax.utils.copy_with_new_vars(mod[gvar]))
+
+        result = (gvar, is_nn_tensor_output)
+        bucket = cls._gvar.setdefault(lookup_key, [])
+        bucket.append((arg_ty, result))
+        return result

python/tvm/relax/frontend/nn/torch.py

@@ -0,0 +1,136 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""PyTorch integration with nn.Module"""
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from tvm_ffi import Array, Shape
+
+from tvm.runtime import Tensor, _tensor
+from tvm.runtime.vm import VirtualMachine
+
+from . import core
+from . import spec as _spec
+
+
+class TorchModule:  # pylint: disable=too-few-public-methods
+    """A wrapper on top of TVM VirtualMachine that takes torch tensors as inputs and returns torch
+    tensors as outputs"""
+
+    spec: _spec.ModuleSpec
+    vm: VirtualMachine  # pylint: disable=invalid-name
+    params: list[Tensor]
+    effects: list[Any]
+
+    def __init__(  # pylint: disable=invalid-name
+        self,
+        spec: _spec.ModuleSpec,
+        vm: VirtualMachine,
+        params: list[Tensor],
+    ):
+        try:
+            self.effects = vm["_initialize_effect"]()
+        except AttributeError:
+            self.effects = None
+
+        self.spec = spec
+        self.vm = vm
+        self.params = params
+
+    def __getitem__(self, method_name: str) -> Callable:
+        def _find_method(method_name):
+            for key, value in zip(self.spec.method_names, self.spec.method_specs):
+                if method_name == key:
+                    return value
+            raise ValueError(f"Method `{method_name}` is not found in the module spec. {self.spec}")
+
+        method_spec = _find_method(method_name)
+        method = self.vm[method_name]
+
+        def _closure(*args):
+            if len(args) != len(method_spec.arg_names):
+                raise TypeError(
+                    f"Argument length mismatch. Expected {len(method_spec.arg_names)} arguments, "
+                    f"but got {len(args)} arguments. The spec is: {method_spec}"
+                )
+            args = [
+                _torch_to_tvm(arg_name, arg_spec, arg)
+                for arg_name, arg_spec, arg in zip(
+                    method_spec.arg_names, method_spec.arg_specs, args
+                )
+            ]
+            if self.effects is not None:
+                outputs, self.effects = method(*args, *self.effects, *self.params)
+            else:
+                outputs = method(*args, *self.params)
+            return _tvm_to_torch(outputs)
+
+        _closure.__name__ = method_name
+        return _closure
+
+
+def _tvm_to_torch(arg):
+    if isinstance(arg, list | tuple | Array):
+        return [_tvm_to_torch(i) for i in arg]
+    if isinstance(arg, _tensor.Tensor):
+        return torch.utils.dlpack.from_dlpack(arg)
+    if isinstance(arg, Shape):
+        return list(arg)
+    raise TypeError(f"Unsupported argument type: {type(arg)}")
+
+
+def _torch_to_tvm(arg_name, arg_spec, arg_torch):
+    if isinstance(arg_spec, _spec.Tensor):
+        if not isinstance(arg_torch, torch.Tensor):
+            raise TypeError(
+                f"Expected argument `{arg_name}` to be `torch.Tensor`, but got {type(arg_torch)}"
+            )
+        return core._from_dlpack(arg_torch)  # pylint: disable=protected-access
+    if isinstance(arg_spec, _spec.Int):
+        if not isinstance(arg_torch, int):
+            raise TypeError(
+                f"Expected argument `{arg_name}` to be `int`, but got {type(arg_torch)}"
+            )
+        return Shape([arg_torch])
+    if isinstance(arg_spec, _spec.Tuple):
+        return [
+            _torch_to_tvm(f"{arg_name}[{i}]", x, arg_torch[i])
+            for i, x in enumerate(arg_spec.elements)
+        ]
+    raise TypeError(f"Unsupported spec item type: {type(arg_spec)}")
+
+
+def _method_spec_from_torch(
+    args_torch: list[Any],
+    method: Callable,
+):
+    def _as_spec(arg_torch):
+        if isinstance(arg_torch, torch.Tensor):
+            _, dtype = str(arg_torch.dtype).rsplit(".", maxsplit=1)
+            return _spec.Tensor(shape=list(arg_torch.shape), dtype=dtype)
+        if isinstance(arg_torch, int):
+            return _spec.Int()
+        raise TypeError(f"Unsupported argument type: {type(arg_torch)}")
+
+    arg_names = list(inspect.signature(method).parameters.keys())
+    if len(arg_names) != len(args_torch):
+        raise TypeError(f"Expected {len(arg_names)} arguments, but got {len(args_torch)} arguments")
+    arg_specs = [_as_spec(i) for i in args_torch]
+    return _spec.MethodSpec(method, arg_names, arg_specs, param_mode="plain", effect_mode="plain")

python/tvm/relax/frontend/nn/visitor.py

@@ -0,0 +1,234 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The visitor and mutator infra for nn.Module."""
+
+from typing import Any
+
+from . import core as nn
+
+
+class Mutator:
+    """The mutator for nn.Module transform. Users can override the `visit_*` methods
+    to apply transform in different structures, or even override the `visit` method
+    to change the logic of traversal."""
+
+    def visit_module(self, name: str, node: nn.Module) -> Any:
+        """The base visiting method for mutation of nn.Module nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Module
+            The current node of nn.Module to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_effect(self, name: str, node: nn.Parameter) -> Any:
+        """The base visiting method for mutation of nn.Parameter nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Parameter
+            The current node of nn.Parameter to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_param(self, name: str, node: nn.Effect) -> Any:
+        """The base visiting method for mutation of nn.Effect nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Effect
+            The current node of nn.Effect to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_moduledict(self, name: str, node: nn.ModuleDict) -> Any:
+        """The base visiting method for mutation of nn.ModuleDict nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ModuleDict
+            The current node of nn.ModuleDict to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_modulelist(self, name: str, node: nn.ModuleList) -> Any:
+        """The base visiting method for mutation of nn.ModuleList nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ModuleList
+            The current node of nn.ModuleList to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_parameterdict(self, name: str, node: nn.ParameterDict) -> Any:
+        """The base visiting method for mutation of nn.ParameterDict nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterDict
+            The current node of nn.ParameterDict to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_parameterlist(self, name: str, node: nn.ParameterList) -> Any:
+        """The base visiting method for mutation of nn.ParameterList nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterList
+            The current node of nn.ParameterList to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit(self, name: str, node: Any) -> Any:
+        """The base dispatching method for visiting of all nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : Any
+            The current node to visit.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+
+        def _get_child_name(parent: str, child: str) -> str:
+            """Get the name of the child node/key given the parent's name."""
+            if parent == "":
+                # in the top level of the module
+                return child
+            else:
+                return f"{parent}.{child}"
+
+        if isinstance(node, nn.ParameterList):
+            for i in range(len(node)):
+                node[i] = self.visit_param(_get_child_name(name, str(i)), node[i])
+        elif isinstance(node, nn.ParameterDict):
+            for k, v in node.items():
+                node[k] = self.visit_param(_get_child_name(name, k), v)
+        elif isinstance(node, nn.ModuleList):
+            for i in range(len(node)):
+                if isinstance(node[i], nn.ParameterDict):
+                    node[i] = self.visit_parameterdict(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ParameterList):
+                    node[i] = self.visit_parameterlist(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ModuleDict):
+                    node[i] = self.visit_moduledict(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.ModuleList):
+                    node[i] = self.visit_modulelist(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Module):
+                    node[i] = self.visit_module(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Effect):
+                    node[i] = self.visit_effect(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Parameter):
+                    node[i] = self.visit_param(f"{name}.{i}", node[i])
+        elif isinstance(node, nn.ModuleDict):
+            for k, v in node.items():
+                if isinstance(v, nn.ParameterDict):
+                    node[k] = self.visit_parameterdict(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ParameterList):
+                    node[k] = self.visit_parameterlist(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ModuleDict):
+                    node[k] = self.visit_moduledict(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ModuleList):
+                    node[k] = self.visit_modulelist(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Module):
+                    node[k] = self.visit_module(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Effect):
+                    node[k] = self.visit_effect(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Parameter):
+                    node[k] = self.visit_param(_get_child_name(name, k), v)
+        else:
+            for key, value in node.__dict__.items():
+                if isinstance(value, nn.ParameterDict):
+                    setattr(node, key, self.visit_parameterdict(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ParameterList):
+                    setattr(node, key, self.visit_parameterlist(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ModuleDict):
+                    setattr(node, key, self.visit_moduledict(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ModuleList):
+                    setattr(node, key, self.visit_modulelist(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Module):
+                    setattr(node, key, self.visit_module(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Effect):
+                    setattr(node, key, self.visit_effect(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Parameter):
+                    setattr(node, key, self.visit_param(_get_child_name(name, key), value))
+        return node

python/tvm/relax/frontend/onnx/__init__.py

@@ -0,0 +1,22 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Tools for converting ONNX graphs into Relax graphs.
+"""
+
+from .onnx_frontend import from_onnx

python/tvm/relax/frontend/stablehlo/__init__.py

@@ -0,0 +1,22 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+StableHLO Frontends for constructing Relax programs, with the model importers
+"""
+
+from .stablehlo_translator import from_stablehlo

python/tvm/relax/frontend/stablehlo/stablehlo_translator.py

@@ -0,0 +1,445 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-outside-toplevel, unused-argument
+
+"""StableHLO frontend of Relax."""
+
+from collections.abc import Callable
+from typing import Any
+
+import tvm
+from tvm import relax, tirx
+
+
+class StableHLOImporter:
+    """An importer from StableHLO to Relax."""
+
+    from jaxlib import mlir
+    from jaxlib.mlir.dialects import stablehlo
+
+    def __init__(self) -> None:
+        from jaxlib import mlir
+
+        self._nodes: dict[str | mlir.ir.Operation, relax.Expr] = {}
+        self.block_builder: relax.BlockBuilder = None
+        self.create_convert_map()
+
+    @staticmethod
+    def _convert_data_type(input_type):
+        """converts the data type from mlir to tvm."""
+        from jaxlib import mlir
+
+        if mlir.ir.ShapedType.isinstance(input_type):
+            input_type = mlir.ir.ShapedType(input_type).element_type
+
+        input_type = str(input_type)
+        if input_type == "f16":
+            return "float16"
+        elif input_type in ["f32", "F32Type"]:
+            return "float32"
+        elif input_type in ["f64", "F64Type"]:
+            return "float64"
+        elif input_type == "i1":
+            return "bool"
+        elif input_type == "i8":
+            return "int8"
+        elif input_type == "i16":
+            return "int16"
+        elif input_type == "i32":
+            return "int32"
+        elif input_type == "i64":
+            return "int64"
+        elif input_type == "ui8":
+            return "uint8"
+        elif input_type == "ui16":
+            return "uint16"
+        elif input_type == "ui32":
+            return "uint32"
+        elif input_type == "ui64":
+            return "uint64"
+        else:
+            raise NotImplementedError(f"input_type {input_type} is not handled yet")
+
+    def _attr2value(self, node) -> Any | list[Any]:
+        import numpy as np
+        from jaxlib import mlir
+
+        if mlir.ir.IntegerAttr.isinstance(node):
+            int_attr = mlir.ir.IntegerAttr(node)
+            return int_attr.value
+        if mlir.ir.FloatAttr.isinstance(node):
+            float_attr = mlir.ir.FloatAttr(node)
+            return float_attr.value
+        if mlir.ir.DenseIntElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseIntElementsAttr(node)
+        elif mlir.ir.DenseFPElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseFPElementsAttr(node)
+        else:
+            raise ValueError("Unsupported Attribute type: " + str(type(node)))
+        ret = []
+        for val in dense_attr:
+            ret.append(val)
+        shape = self.get_shape(node.type)
+        dtype = self._convert_data_type(node.type)
+        return np.asarray(ret, dtype).reshape(shape).tolist()
+
+    def retrieve_operands(self, node):
+        return self._retrieve_operands(node.operands)
+
+    def _retrieve_operands(self, node):
+        from jaxlib import mlir
+
+        # the operand is one of the inputs of FuncOp
+        if isinstance(node, mlir.ir.Operation):
+            return self._nodes[node]
+        if isinstance(node, tuple):
+            return tuple(self._retrieve_operands(x) for x in node)
+        if isinstance(node, list | mlir.ir.OpOperandList):
+            return [self._retrieve_operands(x) for x in node]
+        if isinstance(node, dict):
+            return {self._retrieve_operands(k): self._retrieve_operands(v) for k, v in node.items()}
+        if isinstance(node, mlir.ir.Value):
+            if isinstance(node.owner, mlir.ir.Block):
+                block_arg = mlir.ir.BlockArgument(node)
+                return self._nodes["arg" + str(block_arg.arg_number)]
+            return self._retrieve_operands(node.owner)
+        return node
+
+    def get_shape(self, inpt_type) -> list[Any]:
+        """Get the shape from Type like tensor<?x?xf32>"""
+        from jaxlib import mlir
+
+        shape_type = inpt_type
+        if isinstance(shape_type, mlir.ir.Type):
+            shape_type = mlir.ir.ShapedType(shape_type)
+        ret = []
+        for i in range(shape_type.rank):
+            # get_dim_size
+            if shape_type.is_dynamic_dim(i):
+                n = tirx.Var("n", "int64")
+                ret.append(n)
+            else:
+                ret.append(shape_type.get_dim_size(i))
+
+        return ret
+
+    @staticmethod
+    def _promote_binary_op_args(lhs, rhs):
+        if not isinstance(lhs, relax.Expr) and not isinstance(rhs, relax.Expr):
+            msg = "Both the lhs and the rhs are not expressions."
+            raise AssertionError(msg)
+        if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
+            return lhs, rhs
+        if isinstance(lhs, relax.Expr):
+            assert isinstance(lhs.ty, relax.TensorType)
+            return lhs, relax.const(rhs, lhs.ty.dtype)
+        assert isinstance(rhs.ty, relax.TensorType)
+        return relax.const(lhs, rhs.ty.dtype), rhs
+
+    def _call_binary_op(self, op, lhs, rhs):
+        lhs, rhs = StableHLOImporter._promote_binary_op_args(lhs, rhs)
+        return self.block_builder.emit(op(lhs, rhs))
+
+    def _add(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.add, lhs, rhs)
+        return lhs + rhs
+
+    def _maximum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.maximum(lhs, rhs))
+
+    def _minimum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.minimum(lhs, rhs))
+
+    def _divide(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.divide, lhs, rhs)
+        return lhs / rhs
+
+    def _multiply(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.multiply, lhs, rhs)
+        return lhs * rhs
+
+    def _subtract(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.subtract, lhs, rhs)
+        return lhs - rhs
+
+    def _broadcast_in_dim(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        data = operands[0]
+        # broadcast_dims = self._attr2value(node.attributes["broadcast_dimensions"])
+        shape = self.get_shape(node.result.type)
+        # scalar
+        if len(shape) == 0:
+            return data
+        return self.block_builder.emit(relax.op.broadcast_to(data, shape))
+
+    def _const(self, node: mlir.ir.Operation) -> relax.Expr:
+        const_value = self._attr2value(node.attributes["value"])
+        dtype = self._convert_data_type(node.result.type)
+        return relax.const(const_value, dtype)
+
+    def _dot_general(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.matmul(lhs, rhs))
+
+    def _convolution(self, node) -> relax.Expr:
+        from jaxlib import mlir
+
+        x, weight = self.retrieve_operands(node)
+        shaped_type = mlir.ir.ShapedType(node.result.type)
+        out_dtype = self._convert_data_type(shaped_type.element_type)
+        strides = self._attr2value(node.attributes["window_strides"])
+        padding = self._attr2value(node.attributes["padding"])
+        lhs_dilation = self._attr2value(node.attributes["lhs_dilation"])
+        rhs_dilation = self._attr2value(node.attributes["rhs_dilation"])
+        if len(lhs_dilation) > 0:
+            lhs_dilation = lhs_dilation[0]
+        if len(rhs_dilation) > 0:
+            rhs_dilation = rhs_dilation[0]
+        dilation = (lhs_dilation, rhs_dilation)
+        groups = self._attr2value(node.attributes["batch_group_count"])
+        conv2d = relax.op.nn.conv2d(
+            x,
+            weight,
+            strides=strides,
+            padding=padding[0],
+            dilation=dilation,
+            groups=groups,
+            data_layout="NHWC",
+            kernel_layout="HWIO",
+            out_dtype=out_dtype,
+        )
+
+        return self.block_builder.emit(conv2d)
+
+    def _reshape(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        if isinstance(data, list):
+            assert len(data) == 1
+            data = data[0]
+        new_shape = self.get_shape(node.result.type)
+        return self.block_builder.emit(relax.op.reshape(data, new_shape))
+
+    def _reduce(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        dimensions = self._attr2value(node.attributes["dimensions"])
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.add", f"reducer {reducer_op} in reduce is not supported"
+        return self.block_builder.emit(relax.op.sum(data[0], axis=dimensions))
+
+    def _reduce_window(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        window_dimensions = self._attr2value(node.attributes["window_dimensions"])
+        window_dilations = self._attr2value(node.attributes["window_dilations"])
+
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.maximum", (
+                f"the reducer {reducer_op} in reduce_window is not supported"
+            )
+
+        pool_size = []
+        for i, window_dim in enumerate(window_dimensions):
+            if window_dim == 0:
+                pool_size.append(0)
+            else:
+                dilated_window_size = (window_dim - 1) * window_dilations[i] + 1
+                pool_size.append(dilated_window_size)
+        strides = self._attr2value(node.attributes["window_strides"])
+        # padding = self._attr2value(node.attributes["padding"])
+
+        # TODO (yongwww): Infer the layout automatically
+        layout = "NHWC"
+
+        ret = self.block_builder.emit(
+            relax.op.nn.max_pool2d(
+                operands[0],
+                pool_size=pool_size[1:3],  # HW
+                strides=strides[1:3],
+                padding=[1, 1],
+                dilation=window_dilations[1:3],
+                layout=layout,
+            )
+        )
+        return ret
+
+    def _rsqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.rsqrt(data[0]))
+
+    def _sin(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sin(data[0]))
+
+    def _sinh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sinh(data[0]))
+
+    def _cos(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cos(data[0]))
+
+    def _cosh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cosh(data[0]))
+
+    def _sqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sqrt(data[0]))
+
+    def _round(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.round(data[0]))
+
+    def _exp(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.exp(data[0]))
+
+    def _return(self, node: mlir.ir.Operation) -> relax.Expr:
+        outputs = self.retrieve_operands(node)
+        return self.block_builder.emit_output(self.nodes[outputs])
+
+    def create_convert_map(self):
+        from jaxlib import mlir
+
+        self.convert_map: dict[str, Callable[[mlir.ir.Operation], relax.Var]] = {
+            "stablehlo.add": self._add,
+            "stablehlo.broadcast_in_dim": self._broadcast_in_dim,
+            "stablehlo.constant": self._const,
+            "stablehlo.convolution": self._convolution,
+            "stablehlo.cosine": self._cos,
+            "stablehlo.cosh": self._cosh,
+            "stablehlo.divide": self._divide,
+            "stablehlo.dot_general": self._dot_general,
+            "stablehlo.exponential": self._exp,
+            "stablehlo.maximum": self._maximum,
+            "stablehlo.minimum": self._minimum,
+            "stablehlo.multiply": self._multiply,
+            "stablehlo.reshape": self._reshape,
+            "stablehlo.reduce": self._reduce,
+            "stablehlo.reduce_window": self._reduce_window,
+            "stablehlo.round_nearest_afz": self._round,
+            "stablehlo.rsqrt": self._rsqrt,
+            "stablehlo.sine": self._sin,
+            "chlo.sinh": self._sinh,
+            "stablehlo.sqrt": self._sqrt,
+            "stablehlo.subtract": self._subtract,
+            "func.return": self._return,
+            "stablehlo.return": self._return,
+        }
+
+    def from_stablehlo(self, model, input_info: list[tuple[tuple[int], str]]) -> tvm.IRModule:
+        """Convert a StableHLO Module to a Relax program.
+
+        Parameters
+        ----------
+        model : mlir.ir.Module
+            The StableHLO Module to convert.
+
+        input_info : List[Tuple[Tuple[int], str]]
+            A list of shapes and data types of input tensors.
+
+        Returns
+        -------
+        output : tvm.IRModule
+            The result IRModule with entry function "main"
+        """
+        from jaxlib import mlir
+        from jaxlib.mlir.dialects import stablehlo
+
+        assert isinstance(model, mlir.ir.Module)
+        block: mlir.ir.Block = model.body.operations[0].regions[0].blocks[0]
+
+        # inputs of the function
+        inputs = []
+        for idx, arg in enumerate(block.arguments.types):
+            arg_shape = mlir.ir.ShapedType(arg)
+            ipt_shape = self.get_shape(arg_shape)
+            ipt_dtype = self._convert_data_type(arg_shape.element_type)
+            ipt_name = "arg" + str(idx)
+            ipt_var = relax.Var(f"arg{idx}", relax.TensorType(ipt_shape, ipt_dtype))
+            self._nodes[ipt_name] = ipt_var
+            inputs.append(ipt_var)
+
+        # TODO (yongwww): Handle mlir.ir.Module with multiple functions
+        # Initialize the block builder with a function and a dataflow block.
+        # Raise error if the input stablehlo op is impure
+        func_name = "main"
+        self.block_builder = relax.BlockBuilder()
+
+        with self.block_builder.function(name=func_name, params=inputs.copy()):
+            output = None
+            with self.block_builder.dataflow():
+                block = model.body.operations[0].regions[0].blocks[0]
+                for operation in block.operations:
+                    if isinstance(operation, mlir.dialects.func.ReturnOp | stablehlo.ReturnOp):
+                        operation = operation.operands[0].owner
+                        # TODO (yongwww): handle multiple outputs
+                        output = self.block_builder.emit_output(self._nodes[operation])
+                        break
+
+                    if isinstance(operation, mlir.ir.OpView):
+                        op_name = operation.operation.name
+                        assert op_name in self.convert_map, f"Unsupported operation {op_name}"
+                        self._nodes[operation] = self.convert_map[op_name](operation)
+                    else:
+                        raise ValueError(f"Unsupported op {operation}")
+            assert output is not None
+            self.block_builder.emit_func_output(output)
+
+        mod = self.block_builder.get()
+        return mod
+
+
+def from_stablehlo(
+    stablehlo_module,
+    input_info: list[tuple[tuple[int], str]] | None = None,
+) -> tvm.IRModule:
+    """Convert a StableHLO Module to a Relax program
+
+    Parameters
+    ----------
+    stablehlo_module : Union[str, mlir.ir.Module]
+        The StableHLO Module to convert.
+
+    input_info : List[Tuple[Tuple[int], str]]
+        A list of shapes and data types of input tensors.
+
+    Returns
+    -------
+    output : tvm.IRModule
+        The result IRModule with entry function "main"
+    """
+    from jax._src.interpreters import mlir as jax_mlir
+
+    if isinstance(stablehlo_module, str):
+        # TODO (yongwww): support the serialized bytecode format of StableHLO
+        # model using stablehlo.deserialize_portable_artifact(ir) if the python
+        # binding is ready
+        context = jax_mlir.make_ir_context()
+        stablehlo_module = jax_mlir.ir.Module.parse(stablehlo_module, context)
+    return StableHLOImporter().from_stablehlo(stablehlo_module, input_info)

python/tvm/relax/frontend/tflite/__init__.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Tools for converting TFLite graphs into Relax graphs.
+"""
+
+from .tflite_frontend import from_tflite

python/tvm/relax/frontend/tflite/tflite_flexbuffer.py

@@ -0,0 +1,161 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-argument, too-many-lines, import-outside-toplevel
+# pylint: disable=broad-exception-raised, use-list-literal
+"""Tensorflow lite frontend helper to parse custom options in Flexbuffer format."""
+
+import struct
+from enum import IntEnum
+
+
+class BitWidth(IntEnum):
+    """Flexbuffer bit width schema from flexbuffers.h"""
+
+    BIT_WIDTH_8 = 0
+    BIT_WIDTH_16 = 1
+    BIT_WIDTH_32 = 2
+    BIT_WIDTH_64 = 3
+
+
+class FlexBufferType(IntEnum):
+    """Flexbuffer type schema from flexbuffers.h"""
+
+    FBT_NULL = 0
+    FBT_INT = 1
+    FBT_UINT = 2
+    FBT_FLOAT = 3
+    # Types above stored inline, types below store an offset.
+    FBT_KEY = 4
+    FBT_STRING = 5
+    FBT_INDIRECT_INT = 6
+    FBT_INDIRECT_UINT = 7
+    FBT_INDIRECT_FLOAT = 8
+    FBT_MAP = 9
+    FBT_VECTOR = 10  # Untyped.
+    FBT_VECTOR_INT = 11  # Typed any size (stores no type table).
+    FBT_VECTOR_UINT = 12
+    FBT_VECTOR_FLOAT = 13
+    FBT_VECTOR_KEY = 14
+    FBT_VECTOR_STRING = 15
+    FBT_VECTOR_INT2 = 16  # Typed tuple (no type table, no size field).
+    FBT_VECTOR_UINT2 = 17
+    FBT_VECTOR_FLOAT2 = 18
+    FBT_VECTOR_INT3 = 19  # Typed triple (no type table, no size field).
+    FBT_VECTOR_UINT3 = 20
+    FBT_VECTOR_FLOAT3 = 21
+    FBT_VECTOR_INT4 = 22  # Typed quad (no type table, no size field).
+    FBT_VECTOR_UINT4 = 23
+    FBT_VECTOR_FLOAT4 = 24
+    FBT_BLOB = 25
+    FBT_BOOL = 26
+    FBT_VECTOR_BOOL = 36  # To Allow the same type of conversion of type to vector type
+
+
+class FlexBufferDecoder:
+    """
+    This implements partial flexbuffer deserialization to be able
+    to read custom options. It is not intended to be a general
+    purpose flexbuffer deserializer and as such only supports a
+    limited number of types and assumes the data is a flat map.
+    """
+
+    def __init__(self, buffer):
+        self.buffer = buffer
+
+    def indirect_jump(self, offset, byte_width):
+        """Helper function to read the offset value and jump"""
+        unpack_str = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[byte_width]
+        back_jump = struct.unpack(unpack_str, self.buffer[offset : offset + byte_width])[0]
+        return offset - back_jump
+
+    def decode_keys(self, end, size, byte_width):
+        """Decodes the flexbuffer type vector. Map keys are stored in this form"""
+        # Keys are strings here. The format is all strings separated by null, followed by back
+        # offsets for each of the string. For example, (str1)\0(str1)\0(offset1)(offset2) The end
+        # pointer is pointing at the end of all strings
+        keys = list()
+        for i in range(0, size):
+            offset_pos = end + i * byte_width
+            start_index = self.indirect_jump(offset_pos, byte_width)
+            str_size = self.buffer[start_index:].find(b"\0")
+            assert str_size != -1
+            s = self.buffer[start_index : start_index + str_size].decode("utf-8")
+            keys.append(s)
+        return keys
+
+    def decode_vector(self, end, size, byte_width):
+        """Decodes the flexbuffer vector"""
+        # Each entry in the vector can have different datatype. Each entry is of fixed length. The
+        # format is a sequence of all values followed by a sequence of datatype of all values. For
+        # example - (4)(3.56)(int)(float) The end here points to the start of the values.
+        # Each type byte contains: (type << 2) | bit_width, where bit_width determines actual size.
+        values = list()
+        for i in range(0, size):
+            value_type_pos = end + size * byte_width + i
+            value_type_packed = self.buffer[value_type_pos]
+            value_type = FlexBufferType(value_type_packed >> 2)
+            value_bit_width = BitWidth(value_type_packed & 3)
+            value_byte_width = 1 << value_bit_width
+            value_bytes = self.buffer[
+                end + i * byte_width : end + i * byte_width + value_byte_width
+            ]
+            if value_type == FlexBufferType.FBT_BOOL:
+                value = bool(value_bytes[0])
+            elif value_type == FlexBufferType.FBT_INT:
+                fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            elif value_type == FlexBufferType.FBT_UINT:
+                fmt = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            elif value_type == FlexBufferType.FBT_FLOAT:
+                fmt = {4: "<f", 8: "<d"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            else:
+                raise Exception
+            values.append(value)
+        return values
+
+    def decode_map(self, end, byte_width, parent_byte_width):
+        """Decodes the flexbuffer map and returns a dict"""
+        mid_loc = self.indirect_jump(end, parent_byte_width)
+        size_fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[byte_width]
+        map_size = struct.unpack(size_fmt, self.buffer[mid_loc - byte_width : mid_loc])[0]
+
+        # Find keys
+        keys_offset = mid_loc - byte_width * 3
+        keys_end = self.indirect_jump(keys_offset, byte_width)
+        keys = self.decode_keys(keys_end, map_size, 1)
+
+        # Find values
+        values_end = self.indirect_jump(end, parent_byte_width)
+        values = self.decode_vector(values_end, map_size, byte_width)
+        return dict(zip(keys, values))
+
+    def decode(self):
+        """Decode the buffer. Decoding is partially implemented"""
+        root_end = len(self.buffer) - 1
+        root_byte_width = self.buffer[root_end]
+        root_end -= 1
+        root_packed_type = self.buffer[root_end]
+        root_end -= root_byte_width
+
+        root_type = FlexBufferType(root_packed_type >> 2)
+        byte_width = 1 << BitWidth(root_packed_type & 3)
+
+        if root_type == FlexBufferType.FBT_MAP:
+            return self.decode_map(root_end, byte_width, root_byte_width)
+        raise NotImplementedError("Flexbuffer Decoding is partially imlpemented.")

python/tvm/relax/frontend/torch/__init__.py

@@ -0,0 +1,24 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+PyTorch Frontends for constructing Relax programs, with the model importers
+"""
+
+from .exported_program_translator import from_exported_program
+from .fx_translator import from_fx
+from .dynamo import relax_dynamo, dynamo_capture_subgraphs

python/tvm/relax/frontend/torch/dynamo.py

@@ -0,0 +1,245 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=invalid-name, missing-function-docstring, not-callable
+# pylint: disable=import-outside-toplevel, unused-argument, use-list-literal
+# mypy: ignore-errors
+"""PyTorch Dynamo backend of Relax."""
+
+import functools
+
+import tvm_ffi
+
+import tvm
+from tvm.relax import build as relax_build
+
+from .fx_translator import from_fx
+
+
+def device_from_inputs(example_inputs):
+    for x in example_inputs:
+        if hasattr(x, "device"):
+            return x.device
+    return None
+
+
+def relax_dynamo(pipeline: tvm.transform.Pass | None = None):
+    """A helper function to create a relax backend.
+
+    Parameters
+    ----------
+    pipeline : Optional[tvm.transform.Pass]
+        The pipeline to be applied to the relax module before sent to build.
+
+    Returns
+    -------
+    backend : Callable[[torch.fx.GraphModule, List[torch.Tensor]], Callable]
+        The relax dynamo backend.
+    """
+
+    def _relax_backend(graph_module, example_inputs):
+        import torch  # type: ignore[import]
+
+        assert isinstance(graph_module, torch.fx.GraphModule)
+
+        def to_torch_tensor(nd_tensor):
+            """A helper function to transfer a Tensor to torch.tensor."""
+            if isinstance(nd_tensor, torch.Tensor):
+                # tvm-ffi #517 (Recursive DLPack container conversion) auto-converts
+                # ffi::Tensor items returned in containers back to torch.Tensor when
+                # the call site passed torch.Tensor inputs.
+                return nd_tensor
+            if isinstance(nd_tensor, tvm.runtime.Tensor):
+                return torch.from_numpy(nd_tensor.numpy())
+            elif isinstance(nd_tensor, tvm_ffi.Array):
+                return tuple(to_torch_tensor(x) for x in nd_tensor)
+            else:
+                raise ValueError(f"Unsupported type {type(nd_tensor)}")
+
+        graph_module.graph.eliminate_dead_code()
+
+        device = device_from_inputs(example_inputs)
+
+        assert len(example_inputs)
+
+        fake_inputs = []
+        if isinstance(example_inputs[0], torch._subclasses.fake_tensor.FakeTensor):
+            # Fake tensors
+            fake_inputs = example_inputs
+        else:
+            # Real tensors
+            for node in graph_module.graph.nodes:
+                if node.op != "placeholder":
+                    continue
+                if "grapharg" not in node.meta:
+                    continue
+                fake_tensor = node.meta["grapharg"].fake_tensor
+                if fake_tensor is None:
+                    continue
+                fake_inputs.append(fake_tensor)
+
+        input_info = []
+        shape_vars = {}
+        for tensor in fake_inputs:
+            shape = []
+            for s in tensor.shape:
+                if isinstance(s, torch.SymInt):
+                    if str(s) not in shape_vars:
+                        shape_vars[str(s)] = tvm.tirx.Var(str(s), "int64")
+                    shape.append(shape_vars[str(s)])
+                else:
+                    shape.append(s)
+            input_info.append((shape, tensor.dtype))
+
+        mod = from_fx(graph_module, input_info)
+
+        if device.type == "cuda":
+            dev = tvm.cuda(device.index)
+            target = tvm.target.Target("cuda")
+        else:
+            dev = tvm.cpu(0)
+            target = tvm.target.Target(llvm_target())
+
+        # invoke optimization pipeline.
+        if pipeline is None:
+            # get default pipeline
+            seq = tvm.relax.get_pipeline()
+        elif isinstance(pipeline, str):
+            # lookup by name
+            seq = tvm.relax.get_pipeline(pipeline)
+        else:
+            seq = pipeline
+
+        mod = mod.with_attr("target", target)
+        mod = seq(mod)
+
+        ex = relax_build(mod, target=target)
+
+        vm = tvm.relax.VirtualMachine(ex.mod, device=dev)
+
+        def exec_tvm(*i_args):
+            args = [a.contiguous() for a in i_args if isinstance(a, torch.Tensor)]
+            vm_args = list()
+            for arg in args:
+                if arg.requires_grad:
+                    arg = arg.detach()
+                if isinstance(arg, torch._subclasses.fake_tensor.FakeTensor):
+                    # Materialize a real (eager) Tensor
+                    arg = torch.randn(arg.shape, dtype=arg.dtype, device=device)
+                vm_args.append(arg)
+            outputs = vm["main"](*vm_args)
+            return to_torch_tensor(outputs)
+
+        return exec_tvm
+
+    return _relax_backend
+
+
+def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm.IRModule:
+    """Capture subgraphs of the PyTorch model using torch.compile into an IRModule.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        The PyTorch model to be captured.
+
+    params : List[torch.Tensor]
+        The parameters of the PyTorch model.
+
+    keep_params_as_input : bool
+        Whether to keep model parameters as input variables of the captured Relax functions.
+
+    Returns
+    -------
+    output : ImporterOutput
+        The output of translation, including the translated IRModule.
+        If `keep_params_as_input` is true, the functions in the IRModule have an
+        attribute "params" that contains the weights of the input model. The
+        weights can be detached by `relax.frontend.detach_params`.
+    """
+    import torch  # type: ignore[import]
+    from torch import _dynamo as dynamo  # type: ignore[import]
+    from torch import fx  # type: ignore[import]
+
+    keep_params_as_input = "keep_params_as_input" in kwargs and kwargs["keep_params_as_input"]
+    kwargs.pop("keep_params_as_input", None)
+    mod = tvm.IRModule()
+
+    def _capture(graph_module: fx.GraphModule, example_inputs):
+        assert isinstance(graph_module, torch.fx.GraphModule)
+        input_info = [(tuple(tensor.shape), str(tensor.dtype)) for tensor in example_inputs]
+        mod_ = from_fx(
+            graph_module,
+            input_info,
+            keep_params_as_input=keep_params_as_input,
+            unwrap_unit_return_tuple=True,
+        )
+        new_name = f"subgraph_{len(mod.get_global_vars())}"
+        mod[new_name] = mod_["main"].with_attr("global_symbol", new_name)
+        return graph_module.forward
+
+    dynamo.reset()
+    compiled_model = torch.compile(model, backend=_capture)
+
+    with torch.no_grad():
+        compiled_model(*params, **kwargs)
+
+    return mod
+
+
+@functools.lru_cache(None)
+def llvm_target():
+    import platform
+    import subprocess
+
+    AVX512_TARGET = {"kind": "llvm", "mcpu": "skylake-avx512"}
+    AVX2_TARGET = {"kind": "llvm", "mcpu": "core-avx2"}
+    DEFAULT_TARGET = "llvm"
+
+    system = platform.system()
+
+    if system == "Linux":
+        try:
+            with open("/proc/cpuinfo") as f:
+                cpuinfo = f.read()
+            if "avx512" in cpuinfo:
+                return AVX512_TARGET
+            return AVX2_TARGET
+        except FileNotFoundError:
+            pass
+    elif system == "Darwin":
+        try:
+            result = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.features"],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            if result.returncode == 0:
+                cpu_features = result.stdout.lower()
+                if "avx512" in cpu_features:
+                    return AVX512_TARGET
+                if "avx2" in cpu_features:
+                    return AVX2_TARGET
+        except (FileNotFoundError, subprocess.SubprocessError):
+            pass
+
+        if platform.machine() == "arm64":
+            return DEFAULT_TARGET
+
+    # Default fallback
+    return DEFAULT_TARGET