chore: import upstream snapshot with attribution

python/tvm/relax/backend/cuda/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""The Relax CUDA backend compilation pipeline and other passes."""
+
+from . import flashinfer
+from .pipeline import (
+    finalize_passes,
+    get_default_pipeline,
+    legalize_passes,
+    library_dispatch_passes,
+)

python/tvm/relax/backend/cuda/cublas.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.
+
+"""Pattern table for cuBLAS backend"""
+
+import operator
+from functools import reduce
+
+import tvm
+from tvm import DataType
+from tvm.arith import Analyzer
+from tvm.relax import transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import (
+    make_matmul_dequantize_pattern,
+    make_matmul_multiply_pattern,
+    make_matmul_pattern,
+)
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
+    """Check if dtypes in the given workload are supported by cuBLAS BYOC."""
+    if lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        # The output cannot be 'float8_e5m2' if inputs are 'float8_e4m3fn'
+        return out_dtype != "float8_e5m2"
+    return (
+        (lhs_dtype == "float16" and rhs_dtype == "float16")
+        or (lhs_dtype == "float32" and rhs_dtype == "float32")
+        or (lhs_dtype == "int8" and rhs_dtype == "int8")
+        or (lhs_dtype == "bfloat16" and rhs_dtype == "bfloat16")
+    )
+
+
+def _check_matmul(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    lhs = context.annotated_expr["lhs"]
+    rhs = context.annotated_expr["rhs"]
+    matmul_call = context.annotated_expr["root"]
+
+    if "scale" in context.annotated_expr and "zp" in context.annotated_expr:
+        scale = context.annotated_expr["scale"]
+        zero_point = context.annotated_expr["zp"]
+        # Only scalar values for scale and zero_point are supported.
+        if scale.ty.ndim != 0 or zero_point.ty.ndim != 0:
+            return False
+        # Only zero_point == 0.0 is supported.
+        if zero_point.data.numpy()[()].item() != 0.0:
+            return False
+
+    lhs_dtype = lhs.ty.dtype
+    rhs_dtype = rhs.ty.dtype
+    out_dtype = matmul_call.ty.dtype
+    if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
+        return False
+
+    lhs_shape = lhs.ty.shape.values
+    rhs_shape = rhs.ty.shape.values
+
+    if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+        # Reduction axis must be constant
+        return False
+
+    if lhs_dtype == "int8" and rhs_dtype == "int8":
+        if lhs_shape[-1] % 4 != 0:
+            # Reduction axis must be multiples of 4 for IGEMM
+            return False
+        if not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int) or rhs_shape[-1] % 4 != 0:
+            # Rows number must be multiples of 4 for IGEMM
+            return False
+    elif lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        matmul_rhs_var = matmul_call.args[1]
+        rhs_transposed = False
+        if matmul_rhs_var in context.matched_bindings:
+            matmul_rhs_call = context.matched_bindings[matmul_rhs_var]
+            assert (
+                isinstance(matmul_rhs_call, tvm.relax.Call)
+                and matmul_rhs_call.op.name == "relax.permute_dims"
+            )
+            rhs_transposed = True
+
+        if not rhs_transposed:
+            # cuBLAS FP8 operations require rhs being transposed
+            return False
+
+        # cuBLAS FP8 operations require all tensors being aligned to 16 bytes.
+        if (
+            not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int)
+            or rhs_shape[-1] % (16 // DataType(lhs_dtype).itemsize) != 0
+        ):
+            return False
+        if (
+            not isinstance(rhs_shape[-2], tvm.tirx.expr.IntImm | int)
+            or rhs_shape[-2] % (16 // DataType(out_dtype).itemsize) != 0
+        ):
+            return False
+
+    lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
+    rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
+
+    if "bias" in context.annotated_expr:
+        if lhs_dtype == "int8" and rhs_dtype == "int8":
+            # Non-default epilogue not supported for IGEMM
+            return False
+        bias = context.annotated_expr["bias"]
+        bias_shape = bias.ty.shape.values
+        bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
+        if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
+            # cuBLAS only supports bias vector
+            return False
+
+    analyzer = Analyzer()
+
+    # cuBLASLt does not seem to support batched GEMM with one of matrices having
+    # one batch (with batch_stride 0). So for batched GEMM, the two batch counts
+    # must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
+    # flattening all batch axes into the M axis.
+    return (
+        isinstance(lhs_batches, tvm.tirx.Var)
+        or isinstance(rhs_batches, tvm.tirx.Var)
+        or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
+        or (analyzer.can_prove(lhs_batches >= 1) and analyzer.can_prove(rhs_batches == 1))
+    )
+
+
+register_patterns(
+    [
+        (
+            "cublas.matmul",
+            *make_matmul_pattern(
+                with_bias=False,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed",
+            *make_matmul_pattern(
+                with_bias=False,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_dequantize",
+            *make_matmul_dequantize_pattern(transposed_rhs=True),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_multiply",
+            *make_matmul_multiply_pattern(transposed_rhs=True),
+            _check_matmul,
+        ),
+    ]
+)
+
+
+def partition_for_cublas(mod, bind_constants=False):
+    """
+    Partition the input module into cuBLAS-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    bind_constants : bool
+        Whether or not to keep bound constants in the grouped function.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the cuBLAS backend.
+    """
+
+    patterns = get_patterns_with_prefix("cublas")
+    return transform.FuseOpsByPattern(
+        patterns, bind_constants=bind_constants, annotate_codegen=True
+    )(mod)

python/tvm/relax/backend/cuda/cudnn.py

@@ -0,0 +1,202 @@
+# 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.
+
+"""Pattern table for cuDNN backend"""
+
+import operator
+from functools import partial, reduce
+
+import tvm
+from tvm import relax
+from tvm.relax import PyExprMutator, expr_functor, transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_conv2d_pattern, make_stacked_attention_pattern
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+    """Check if dtypes in the given workload are supported by cuDNN BYOC."""
+    return (lhs_dtype == "float16" and rhs_dtype == "float16") or (
+        lhs_dtype == "float32" and rhs_dtype == "float32"
+    )
+
+
+def _is_supported_format(data_layout, kernel_layout):
+    """Check if layouts in the given workload are supported by cuDNN BYOC."""
+    return (data_layout == "NHWC" and kernel_layout == "OHWI") or (
+        data_layout == "NCHW" and kernel_layout == "OIHW"
+    )
+
+
+def _check_conv2d(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    # Retrieve the annotated expression from context
+    conv2d_call = context.annotated_expr["root"]
+    input_expr = context.annotated_expr["input"]
+    weight_expr = context.annotated_expr["weight"]
+
+    # Check if the data types of input and weights are supported by cuDNN BYOC
+    input_dtype = input_expr.ty.dtype
+    weight_dtype = weight_expr.ty.dtype
+    if not _is_supported_dtype(input_dtype, weight_dtype):
+        return False
+
+    input_layout = conv2d_call.attrs.data_layout
+    weight_layout = conv2d_call.attrs.kernel_layout
+    if not _is_supported_format(input_layout, weight_layout):
+        return False
+
+    return True
+
+
+def _check_stacked_attention(context: PatternCheckContext, layout: str) -> bool:
+    """Check if the given stacked attention workload can be offloaded to cuDNN."""
+    if has_leaking_intermediate_variables(context):
+        return False
+    if layout == "BS3NH":
+        if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
+            return False
+        if "split" in context.annotated_expr:
+            split_op = context.annotated_expr["split"]
+            if not split_op.attrs.axis == 2:
+                return False
+    elif layout == "SBN3H":
+        if not context.annotated_expr["stacked_qkv"].ty.ndim == 4:
+            return False
+        if "split" in context.annotated_expr:
+            split_op = context.annotated_expr["split"]
+            if not split_op.attrs.axis == 3:
+                return False
+    else:
+        raise NotImplementedError(f"Unsupported layout: {layout}")
+    return True
+
+
+register_patterns(
+    [
+        (
+            "cudnn.conv2d.nhwc_ohwi",
+            *make_conv2d_pattern(
+                with_bias=False,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias",
+            *make_conv2d_pattern(
+                with_bias=True,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias_relu",
+            *make_conv2d_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.attention.BS3NH",
+            *make_stacked_attention_pattern(start_op="split", layout="BS3NH"),
+            partial(_check_stacked_attention, layout="BS3NH"),
+        ),
+        (
+            "cudnn.attention.SBN3H",
+            *make_stacked_attention_pattern(start_op="split", layout="SBN3H"),
+            partial(_check_stacked_attention, layout="SBN3H"),
+        ),
+    ]
+)
+
+
+def partition_for_cudnn(mod):
+    """
+    Partition the input module into cuDNN-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the cuDNN backend.
+    """
+
+    patterns = get_patterns_with_prefix("cudnn")
+    return tvm.transform.Sequential(
+        [
+            transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True),
+            annotate_workspace,
+            transform.AllocateWorkspace(),
+        ]
+    )(mod)
+
+
+def _shape_1d(shape):
+    return reduce(operator.mul, shape, 1)
+
+
+@expr_functor.mutator
+class WorkspaceAnnotator(PyExprMutator):
+    """Annotate a workspace requirement for each cuDNN-offloaded function."""
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_function_(self, f):
+        if "Composite" not in f.attrs:
+            body = super().visit_expr(f.body)
+            new_f = relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+            if "global_symbol" in f.attrs and "cudnn" in f.attrs["global_symbol"]:
+                composite_func = body.blocks[0].bindings[0].value
+                if "WorkspaceSize" in composite_func.attrs:
+                    return new_f.with_attr("WorkspaceSize", composite_func.attrs["WorkspaceSize"])
+
+            return new_f
+
+        if "attention" in f.attrs["Composite"] and "cudnn" in f.attrs["Composite"]:
+            # Workspace is needed only for larger head sizes, but for simplicity we always allocate.
+            out_dtype = f.ret_ty.dtype
+            out_size_1d = _shape_1d(f.ret_ty.shape)
+            # This needs to be in sync with the actual value that the kernel expects.
+            workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
+            if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+                # Tempororay workaround for dynamic shape workload. Will be removed when
+                # workspace for dynamic shape workload is implemented.
+                workspace_size_bytes = 8
+            return f.with_attr("WorkspaceSize", workspace_size_bytes)
+
+        return f
+
+
+@tvm.transform.module_pass(opt_level=0)
+def annotate_workspace(mod, _):
+    """Pass to annotate a workspace requirement for each cuDNN-offloaded function."""
+    annotator = WorkspaceAnnotator(mod)
+    for name, f in mod.functions_items():
+        if isinstance(f, relax.Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
+    return mod

python/tvm/relax/backend/cuda/cutlass.py

@@ -0,0 +1,620 @@
+# 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
+# ruff: noqa: E731
+"""Pattern table for CUTLASS backend"""
+
+import operator
+from collections.abc import Mapping, Sequence
+from functools import reduce
+
+import tvm
+from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
+from tvm.relax import (
+    Call,
+    ExternFunc,
+    Function,
+    PyExprMutator,
+    Var,
+    expr_functor,
+    transform,
+)
+from tvm.relax.dpl import rewrite_call
+from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import (
+    make_attention_pattern,
+    make_attention_rewrite_pattern,
+    make_fused_bias_activation_pattern,
+    make_layer_norm_pattern,
+    make_matmul_pattern,
+    make_residual_block_pattern,
+    make_rms_norm_pattern,
+    make_stacked_attention_pattern,
+)
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+    """Check if dtypes in the given workload are supported by CUTLASS."""
+    return (
+        (lhs_dtype == "float16" and rhs_dtype == "float16")
+        or (lhs_dtype == "float32" and rhs_dtype == "float32")
+        or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+    )
+
+
+def _shape_1d(shape):
+    return reduce(operator.mul, shape, 1)
+
+
+def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequence[Var]]):
+    if from_var == to_var:
+        return True
+
+    checked = set()
+    vars_to_check = [to_var]
+    while vars_to_check:
+        current_var = vars_to_check.pop()
+        for user in var_usages.get(current_var, []):
+            if user == from_var:
+                return True
+            if user not in checked:
+                checked.add(user)
+                vars_to_check.append(user)
+
+    return False
+
+
+def _is_same_shape(shape1, shape2):
+    analyzer = tvm.arith.Analyzer()
+    return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)])
+
+
+def _is_bias_like(shape, out_channel):
+    return shape[-1] == out_channel and _shape_1d(shape) == out_channel
+
+
+def _check_residual(root_call: Call, context: PatternCheckContext) -> bool:
+    if "residual" in context.annotated_expr:
+        residual = context.annotated_expr["residual"]
+        if not isinstance(residual, Var):
+            if residual not in context.value_to_bound_var:
+                return False
+
+            residual = context.value_to_bound_var[residual]
+
+        root_var = context.value_to_bound_var[root_call]
+        if _has_dependency(from_var=residual, to_var=root_var, var_usages=context.var_usages):
+            # If residual depends on the result of the root call, this cannot be handled by cutlass.
+            return False
+
+        shape1 = root_var.ty.shape
+        shape2 = residual.ty.shape
+        out_channel = shape1[-1]
+
+        if not _is_same_shape(shape1, shape2) and not _is_bias_like(shape2, out_channel):
+            return False
+
+    return True
+
+
+def _check_conv2d(context: PatternCheckContext) -> bool:
+    """Check if the given conv2d workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+
+    conv2d_call = context.annotated_expr["root"]
+    data_layout = conv2d_call.attrs.data_layout
+    kernel_layout = conv2d_call.attrs.kernel_layout
+    data, weight, *_ = conv2d_call.args
+    if (
+        data_layout != "NHWC"
+        or kernel_layout != "OHWI"
+        or not _is_supported_dtype(data.ty.dtype, weight.ty.dtype)
+    ):
+        return False
+
+    if not _check_residual(conv2d_call, context):
+        return False
+
+    # Check if any dimensions are symbolic.
+    for dim in data.ty.shape.values:
+        if isinstance(dim, tvm.tirx.Var):
+            return False
+
+    # pylint: disable=invalid-name
+    IC = data.ty.shape.values[3]
+    OC = weight.ty.shape.values[0]
+    # not depthwise conv2d
+    return not IC == OC == conv2d_call.attrs.groups
+
+
+def _check_matmul(context: PatternCheckContext) -> bool:
+    """Check if the given matmul workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+
+    lhs = context.annotated_expr["lhs"]
+    rhs = context.annotated_expr["rhs"]
+
+    lhs_dtype = lhs.ty.dtype
+    rhs_dtype = rhs.ty.dtype
+    if not _is_supported_dtype(lhs_dtype, rhs_dtype):
+        return False
+
+    if not _check_residual(context.annotated_expr["root"], context):
+        return False
+
+    lhs_shape = lhs.ty.shape.values
+    rhs_shape = rhs.ty.shape.values
+    return is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape)
+
+
+def _get_activation_from_name(pattern_name):
+    if "_relu" in pattern_name:
+        return "relax.nn.relu"
+    elif "_gelu_tanh" in pattern_name:
+        return "relax.nn.gelu_tanh"
+    elif "_gelu" in pattern_name:
+        return "relax.nn.gelu"
+    elif "_silu" in pattern_name:
+        return "relax.nn.silu"
+    else:
+        return None
+
+
+def matmul_patterns():
+    """
+    Returns a list of all matmul patterns in cutlass BYOC backend.
+    """
+
+    def _matmul_pattern(pattern_name):
+        transposed_rhs = "_transposed" in pattern_name
+        with_bias = "_bias" in pattern_name
+        activation = _get_activation_from_name(pattern_name)
+
+        return (
+            pattern_name,
+            *make_matmul_pattern(
+                transposed_rhs=transposed_rhs,
+                with_bias=with_bias,
+                activation=activation,
+            ),
+            _check_matmul,
+        )
+
+    return [
+        _matmul_pattern("cutlass.matmul"),
+        _matmul_pattern("cutlass.matmul_bias"),
+        _matmul_pattern("cutlass.matmul_bias_relu"),
+        _matmul_pattern("cutlass.matmul_bias_gelu"),
+        _matmul_pattern("cutlass.matmul_transposed"),
+        _matmul_pattern("cutlass.matmul_transposed_bias"),
+        _matmul_pattern("cutlass.matmul_transposed_bias_relu"),
+        _matmul_pattern("cutlass.matmul_transposed_bias_gelu"),
+    ]
+
+
+def _check_decode_matmul(ctx):
+    """Check if the given decode -> matmul workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(ctx):
+        return False
+
+    root = ctx.annotated_expr["root"]
+
+    if not _check_residual(root, ctx):
+        return False
+
+    # out_dtype = "float32" not supported unless matmul is followed by cast to fp16.
+    if root.ty.dtype == "float32":
+        return False
+
+    call_tir_decode = ctx.annotated_expr["w_decoded"]
+    if "decode" not in call_tir_decode.args[0].name_hint:
+        return False
+
+    N = root.ty.shape[-1]
+
+    if ctx.annotated_expr["lhs"].ty.dtype != "float16":
+        return False
+
+    # weight needs to be packed to int8.
+    packed_weight = ctx.annotated_expr["w_encoded"]
+
+    if packed_weight.ty.dtype != "int8":
+        return False
+
+    # The kernel expects the weight to be preprocessed by this packed function.
+    if (
+        isinstance(packed_weight, Call)
+        and isinstance(packed_weight.args[0], ExternFunc)
+        and packed_weight.args[0].global_symbol != "cutlass.ft_preprocess_weight"
+    ):
+        return False
+
+    scales = ctx.annotated_expr["scales"]
+
+    if scales.ty.dtype != "float16":
+        return False
+
+    # scale shape needs to be (N,) or (1, N) or (K // group_size, N)
+    if len(scales.ty.shape) > 2 or scales.ty.shape[-1] != N:
+        return False
+
+    if "bias" in ctx.annotated_expr:
+        out_shape = root.ty.shape
+        bias_shape = ctx.annotated_expr["bias"].ty.shape
+
+        # bias shape needs to be (N,), possibly with additional axes on the front.
+        # It can also have the same shape as the output.
+        if not _is_bias_like(bias_shape, N) and not _is_same_shape(out_shape, bias_shape):
+            return False
+
+    return True
+
+
+def decode_matmul_patterns():
+    """Returns a list of supported decode -> matmul patterns."""
+
+    def _decode_matmul_pattern(name):
+        scales = wildcard()
+        x = wildcard()
+        w_packed = wildcard()
+
+        w = is_op("relax.call_tir")(
+            GlobalVarPattern(),
+            TuplePattern([w_packed, scales]),
+        )
+        matmul = is_op("relax.matmul")(x, w)
+
+        if "cast" in name:
+            matmul = is_op("relax.astype")(matmul)
+
+        annotations = {
+            "root": matmul,
+            "lhs": x,
+            "w_encoded": w_packed,
+            "w_decoded": w,
+            "scales": scales,
+        }
+
+        if "bias" in name:
+            annotations["bias"] = bias = wildcard()
+            out = is_op("relax.add")(matmul, bias)
+        else:
+            out = matmul
+
+        if "gelu" in name:
+            out = is_op("relax.nn.gelu")(out)
+
+        return name, out, annotations, _check_decode_matmul
+
+    return [
+        _decode_matmul_pattern("cutlass.decode_matmul"),
+        _decode_matmul_pattern("cutlass.decode_matmul_bias"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias"),
+        _decode_matmul_pattern("cutlass.decode_matmul_bias_gelu"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias_gelu"),
+    ]
+
+
+def conv2d_patterns():
+    """
+    Returns a list of all conv2d patterns in cutlass BYOC backend.
+    """
+
+    def _conv2d_pattern(pattern_name):
+        with_bias = "_bias" in pattern_name
+        activation = _get_activation_from_name(pattern_name)
+
+        return (
+            pattern_name,
+            *make_fused_bias_activation_pattern(
+                "relax.nn.conv2d",
+                with_bias=with_bias,
+                activation=activation,
+            ),
+            _check_conv2d,
+        )
+
+    return [
+        _conv2d_pattern("cutlass.conv2d"),
+        _conv2d_pattern("cutlass.conv2d_bias"),
+        _conv2d_pattern("cutlass.conv2d_bias_relu"),
+        _conv2d_pattern("cutlass.conv2d_bias_silu"),
+    ]
+
+
+def residual_block_patterns():
+    """
+    Returns a list of all residual block patterns in cutlass BYOC backend.
+    """
+    patterns = []
+
+    for activation, name_postfix in [(None, ""), ("relax.nn.relu", "_relu")]:
+        for check, base_patterns in [
+            (_check_conv2d, conv2d_patterns()),
+            (_check_matmul, matmul_patterns()),
+            (_check_decode_matmul, decode_matmul_patterns()),
+        ]:
+            for name, pat, arg_pat, _ in base_patterns:
+                # Append residual patterns only to those base patterns with bias add,
+                # since conv2d or matmul + residual add without bias is already supported
+                # via conv2d or matmul + bias patterns (the residual input is treated as "bias").
+                if "bias" in name:
+                    for bin_op in ["relax.add", "relax.multiply"]:
+                        patterns.append(
+                            (
+                                name + "_residual_" + bin_op.split(".")[-1] + name_postfix,
+                                *make_residual_block_pattern(
+                                    (pat, arg_pat), binary_op=bin_op, activation=activation
+                                ),
+                                check,
+                            )
+                        )
+
+    return patterns
+
+
+def _check_stacked_attention(context: PatternCheckContext) -> bool:
+    """Check if the given stacked attention workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+    if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
+        return False
+    if "split" in context.annotated_expr:
+        split_op = context.annotated_expr["split"]
+        if not split_op.attrs.axis == 2:
+            return False
+    else:
+        get_const_int_list = lambda tup: [int(e.value) for e in tup]
+        last_end = 0
+        for name in ["query", "key", "value"]:
+            assert f"strided_slice_{name}" in context.annotated_expr
+            strided_slice_op = context.annotated_expr[f"strided_slice_{name}"]
+            axes = get_const_int_list(strided_slice_op.args[1])
+            begins = get_const_int_list(strided_slice_op.args[2])
+            ends = get_const_int_list(strided_slice_op.args[3])
+            strides = get_const_int_list(strided_slice_op.args[4])
+
+            if axes != [2]:
+                return False
+            if begins != [last_end]:
+                return False
+            if not len(ends) == 1:
+                return False
+            if strides != [1]:
+                return False
+            last_end = ends[0]
+    return True
+
+
+def attention_patterns():
+    """
+    Returns a list of all attention patterns in cutlass BYOC backend.
+    """
+    return [
+        (
+            "cutlass.attention",
+            *make_attention_pattern(),
+        ),
+        (
+            "cutlass.attention_bias",
+            *make_attention_pattern(with_bias=True),
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="split"),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="split", with_bias=True),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="strided_slice"),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="strided_slice", with_bias=True),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.attention_var_len",
+            *make_attention_pattern(var_len=True),
+        ),
+    ]
+
+
+def _check_layer_norm(context: PatternCheckContext) -> bool:
+    attrs = context.matched_expr.attrs
+
+    if not attrs.center or not attrs.scale:
+        return False
+
+    if len(attrs.axes) != 1:
+        # Contiguous inner-most axes can be supported, but reject it for now for simplicity.
+        return False
+
+    axis = int(attrs.axes[0])
+    rank = len(context.matched_expr.ty.shape)
+
+    if axis < 0:
+        axis += rank
+
+    return axis == rank - 1
+
+
+def layer_norm_pattern():
+    """Create a layer norm pattern for CUTLASS."""
+    return [
+        (
+            "cutlass.layer_norm",
+            *make_layer_norm_pattern(),
+            _check_layer_norm,
+        ),
+    ]
+
+
+def _check_rms_norm(ctx: PatternCheckContext) -> bool:
+    rms_norm = ctx.annotated_expr["rms_norm"]
+    if "rms_norm" not in rms_norm.args[0].name_hint:
+        return False
+
+    return True
+
+
+def rms_norm_pattern():
+    """Create a RMS norm pattern for CUTLASS."""
+    return [
+        (
+            "cutlass.rms_norm",
+            *make_rms_norm_pattern(),
+            _check_rms_norm,
+        ),
+    ]
+
+
+def attention_rewrite_patterns():
+    """
+    Returns a list of all attention rewriting patterns in cutlass BYOC backend.
+    """
+    patterns = []
+    for qkv_layout in ["BSNH", "BSH"]:
+        for out_layout in ["BSNH", "BSH"]:
+            for with_bias in [True, False]:
+                for with_cast in [True, False]:
+                    patterns.append(
+                        make_attention_rewrite_pattern(qkv_layout, out_layout, with_bias, with_cast)
+                    )
+    return patterns
+
+
+register_patterns(
+    [
+        *conv2d_patterns(),
+        *matmul_patterns(),
+        *decode_matmul_patterns(),
+        *residual_block_patterns(),
+        *attention_patterns(),
+        *layer_norm_pattern(),
+        *rms_norm_pattern(),
+    ]
+)
+
+_REWRITE_PATTERNS = [*attention_rewrite_patterns()]
+
+
+@expr_functor.mutator
+class WorkspaceAnnotator(PyExprMutator):
+    """Annotate a workspace requirement for each CUTLASS-offloaded function."""
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_function_(self, f):
+        if "Composite" not in f.attrs:
+            body = super().visit_expr(f.body)
+            new_f = Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+            if "global_symbol" in f.attrs and "cutlass" in f.attrs["global_symbol"]:
+                composite_func = body.blocks[0].bindings[0].value
+                if "WorkspaceSize" in composite_func.attrs:
+                    return new_f.with_attr("WorkspaceSize", composite_func.attrs["WorkspaceSize"])
+
+            return new_f
+
+        if "attention" in f.attrs["Composite"] and "cutlass" in f.attrs["Composite"]:
+            # Workspace is needed only for larger head sizes, but for simplicity we always allocate.
+            out_dtype = f.ret_ty.dtype
+            out_size_1d = _shape_1d(f.ret_ty.shape)
+            # This needs to be in sync with the actual value that the kernel expects.
+            workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
+            if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+                # Tempororay workaround for dynamic shape workload. Will be removed when
+                # workspace for dynamic shape workload is implemented.
+                workspace_size_bytes = 8
+            return f.with_attr("WorkspaceSize", workspace_size_bytes)
+
+        return f
+
+
+@tvm.transform.module_pass(opt_level=0)
+def annotate_workspace(mod, _):
+    """Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
+    annotator = WorkspaceAnnotator(mod)
+    for name, f in mod.functions_items():
+        if isinstance(f, Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
+    return mod
+
+
+def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
+    """
+    Partition the input module into CUTLASS-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    annotate_codegen: bool
+        Whether to wrap each created composite function with another function, whose
+        body consists only of a call to the composite function. See the doc of FuseOpsByPattern
+        for more detail.
+
+    use_flash_mqa: bool
+        Whether to consider a rewrite pattern for multi-query attention, which is supported by
+        the Flash Attention kernel.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        compiled by the CUTLASS backend.
+    """
+    for func_name, func in mod.functions_items():
+        if isinstance(func, Function):
+            if use_flash_mqa:
+                mqa_pattern, rewriter = make_attention_rewrite_pattern(
+                    "BSNH", "BSNH", with_bias=False, with_cast=True, with_kv_repeat=True
+                )
+                func = rewrite_call(mqa_pattern, rewriter, func)
+
+            for pattern, rewriter in _REWRITE_PATTERNS:
+                func = rewrite_call(pattern, rewriter, func)
+
+        mod[func_name] = func
+
+    patterns = get_patterns_with_prefix("cutlass")
+    return tvm.transform.Sequential(
+        [
+            transform.FuseOpsByPattern(
+                patterns, bind_constants=False, annotate_codegen=annotate_codegen
+            ),
+            annotate_workspace,
+            transform.AllocateWorkspace(),
+        ]
+    )(mod)

python/tvm/relax/backend/cuda/flashinfer.py

@@ -0,0 +1,345 @@
+# 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.
+
+"""FlashInfer JIT compilation module for CUDA backend"""
+
+import re
+from pathlib import Path
+
+import tvm
+from tvm.target import Target
+
+
+def _rename_exported_func_names(source_paths: list[Path], prefix: str):
+    """Rename the ffi-exported function names in the source files to the given prefix."""
+    pattern = re.compile(r"^(\s*TVM_FFI_DLL_EXPORT_TYPED_FUNC\()([A-Za-z0-9_]+)(,.*)$")
+    for source_path in source_paths:
+        if not source_path.name.endswith("_binding.cu"):
+            continue
+
+        original_text = source_path.read_text(encoding="utf-8")
+        lines = original_text.splitlines(keepends=True)
+        updated = False
+        for idx, line in enumerate(lines):
+            line_body = line.rstrip("\r\n")
+            line_ending = line[len(line_body) :]
+            match = pattern.match(line_body)
+            if not match:
+                continue
+            new_body = f"{match.group(1)}{prefix}_{match.group(2)}{match.group(3)}"
+            lines[idx] = new_body + line_ending
+            updated = True
+
+        if updated:
+            source_path.write_text("".join(lines), encoding="utf-8")
+
+
+def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm.runtime.Module]:
+    return [
+        tvm.runtime.load_static_library(str(obj_path.absolute()), func_names=[])
+        for obj_path in object_files
+    ]
+
+
+def gen_flashinfer_prefill_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    qk_head_dim: int,
+    v_head_dim: int,
+    enable_inline_rope: bool,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for prefill.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    qk_head_dim : int
+        The head dimension of the query and key tensors.
+    v_head_dim : int
+        The head dimension of the value tensor.
+    enable_inline_rope : bool
+        Whether to enable inline rotary positional embedding.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer prefill kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_customize_batch_prefill_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    if enable_inline_rope and qk_head_dim != v_head_dim:
+        raise ValueError("Inline rope mode is not supported when qk_head_dim == v_head_dim")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    # Todo(tvm-team): decide which backend ("fa2/fa3") to use
+    backend = "fa2"
+    variant_name = (
+        "DefaultAttention<false, false, false, false>"
+        if backend == "fa2"
+        else "DefaultAttention<false>"
+    )
+    variant_decl = (
+        "#include <flashinfer/attention/variants.cuh>"
+        if backend == "fa2"
+        else "#include <flashinfer/attention/hopper/variants.cuh>"
+    )
+    jit_spec = gen_customize_batch_prefill_module(
+        backend=backend,
+        uri=f"batch_prefill_tvm_dtype_q_{dtype_q}_"
+        + f"dtype_kv_{dtype_kv}_"
+        + f"dtype_o_{dtype_o}_"
+        + f"qk_head_dim_{qk_head_dim}_"
+        + f"v_head_dim_{v_head_dim}_"
+        + f"enable_inline_rope_{enable_inline_rope}",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        idtype=torch.int32,
+        head_dim_qk=qk_head_dim,
+        head_dim_vo=v_head_dim,
+        pos_encoding_mode=int(enable_inline_rope),
+        additional_tensor_names=[],
+        additional_tensor_dtypes=[],
+        additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
+        additional_scalar_dtypes=["double", "double", "double"],
+        variant_name=variant_name,
+        variant_decl=variant_decl,
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_prefill")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_flashinfer_decode_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    qk_head_dim: int,
+    v_head_dim: int,
+    enable_inline_rope: bool,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for decode.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    qk_head_dim : int
+        The head dimension of the query and key tensors.
+    v_head_dim : int
+        The head dimension of the value tensor.
+    enable_inline_rope : bool
+        Whether to enable inline rotary positional embedding.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer decode kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_customize_batch_decode_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    jit_spec = gen_customize_batch_decode_module(
+        uri=f"batch_decode_tvm_dtype_q_{dtype_q}_"
+        + f"dtype_kv_{dtype_kv}_"
+        + f"dtype_o_{dtype_o}_"
+        + f"qk_head_dim_{qk_head_dim}_"
+        + f"v_head_dim_{v_head_dim}_"
+        + f"enable_inline_rope_{enable_inline_rope}",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        idtype=torch.int32,
+        head_dim_qk=qk_head_dim,
+        head_dim_vo=v_head_dim,
+        pos_encoding_mode=int(enable_inline_rope),
+        additional_tensor_names=[],
+        additional_tensor_dtypes=[],
+        additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
+        additional_scalar_dtypes=["double", "double", "double"],
+        variant_name="DefaultAttention<false, false, false, false>",
+        variant_decl="#include <flashinfer/attention/variants.cuh>",
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_decode")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_flashinfer_mla_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    head_dim_ckv: int,
+    head_dim_kpe: int,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for MLA.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    head_dim_ckv : int
+        The head dimension of the compressed key/value tensors.
+    head_dim_kpe : int
+        The head dimension of the query/key positional embedding.
+    target : Target
+        The target device to compile for.
+    num_threads : int
+        The number of threads to use for compilation.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer MLA kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_batch_mla_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    jit_spec = gen_batch_mla_module(
+        backend="fa2",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        dtype_idx=torch.int32,
+        head_dim_ckv=head_dim_ckv,
+        head_dim_kpe=head_dim_kpe,
+        use_profiler=False,
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_mla")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_grouped_gemm_module(
+    target: Target, return_static_libs: bool = False
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for FP8 grouped GEMM.
+
+    Parameters
+    ----------
+    target : Target
+        The target device to compile for.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    List[tvm.runtime.Module]
+        A list of compiled static library modules for FlashInfer FP8 grouped GEMM kernels.
+
+    Note
+    _____
+    when apply grouped gemm on A: (total_m, k), B: (batch_size, n, k), m_indptr: (batch_size, )
+    requires all m in m_indptr to be multiple of 4
+    """
+    # NOTE: This function is still under development,
+    # and we currently only support SM100 grouped gemm
+    try:
+        from flashinfer.gemm import (  # pylint: disable=import-outside-toplevel
+            gen_gemm_sm100_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+
+    compute_version = "".join(tvm.support.nvcc.get_target_compute_version(target).split("."))
+    if compute_version == "100":
+        jit_spec = gen_gemm_sm100_module()
+    else:
+        raise ValueError(f"Unsupported compute version: {compute_version}")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]

python/tvm/relax/backend/cuda/pipeline.py

@@ -0,0 +1,90 @@
+# 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 Relax CUDA backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for CUDA backend."""
+    return [
+        relax.backend.DispatchSampling(),
+        relax.backend.DispatchSortScan(),
+    ]
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for CUDA backend."""
+    from tvm.s_tir import dlight as dl  # pylint: disable=import-outside-toplevel
+
+    return [
+        tvm.relax.transform.LegalizeOps(),
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FoldConstant(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+        dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        ),
+    ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for CUDA backend."""
+    return [
+        relax.transform.RewriteDataflowReshape(),
+        relax.transform.ToNonDataflow(),
+        relax.transform.RemovePurityChecking(),
+        relax.transform.CallTIRRewrite(),
+    ]
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for CUDA backend."""
+    return [
+        relax.transform.StaticPlanBlockMemory(),
+        relax.transform.RewriteCUDAGraph(),
+        relax.transform.LowerAllocTensor(),
+        relax.transform.KillAfterLastUse(),
+        relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.ComputePrimValue(),
+        relax.transform.VMShapeLower(),
+        relax.transform.AttachGlobalSymbol(),
+    ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for CUDA."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline