Note
This page is a rendered version of 002_bring_your_own_kernel.ipynb available on GitHub.
Adding new Operators to CUTLASS Operator API#
CUTLASS Operator API is designed to make it easy for users to add their own kernels so that it can be discovered and run as Operators under the uniform API. We welcome contributions of kernels written in CUTLASS Python by “bringing your own operator.”
This example shows how to add a CuTe DSL operator to the CUTLASS Operator API.
Bring your own implementation#
Kernel authors wishing to add a CuTe DSL operator to Operator API likely already have the kernel written in CuTe DSL, but have not yet implemented the API’s needed Operator interface. Within the API, we separate these components into the “kernel” or “implementation” itself, which is written in CuTe DSL, and the Operator that implements API methods to expose the underlying kernel via a uniform interface.
For example, consider the following implementation of a simple FP64 GEMM kernel:
[1]:
from collections.abc import Callable
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
class F64GemmKernelImplementation:
def __init__(self, cta_tile_shape_mn: tuple[int, int]):
self.cta_tile_shape_mn = cta_tile_shape_mn
@cute.jit
def __call__(
self, a: cute.Tensor, b: cute.Tensor, out: cute.Tensor, stream: cuda.CUstream
):
l, m, n = out.shape # noqa: E741 (L batch dim)
m_tiles = (m + self.cta_tile_shape_mn[0] - 1) // self.cta_tile_shape_mn[0]
n_tiles = (n + self.cta_tile_shape_mn[1] - 1) // self.cta_tile_shape_mn[1]
grid = (m_tiles, n_tiles, l)
block = [self.cta_tile_shape_mn[0], self.cta_tile_shape_mn[1], 1]
self.kernel(a, b, out).launch(grid=grid, block=block, stream=stream)
@cute.kernel
def kernel(self, a: cute.Tensor, b: cute.Tensor, out: cute.Tensor):
_L, M, N = out.shape
K = a.shape[-1]
m_tile, n_tile, l_idx = cute.arch.block_idx()
tidx, tidy, _ = cute.arch.thread_idx()
m_idx = m_tile * self.cta_tile_shape_mn[0] + tidx
n_idx = n_tile * self.cta_tile_shape_mn[1] + tidy
if m_idx < M and n_idx < N:
out[l_idx, m_idx, n_idx] = cutlass.Float64(0)
for k_idx in range(K):
out[l_idx, m_idx, n_idx] += (
a[l_idx, m_idx, k_idx] * b[l_idx, k_idx, n_idx]
)
The implementation is configurable via a cta_tile_shape_mn argument, which controls the size of blocks and tiles in the M and N modes. A simple cute.jit function computes the grid and block size for the input problem based on cta_tile_shape_mn, and launches the kernel. The cute.kernel itself simply has each thread compute a single output element of the matrix by taking a dot product.
This implementation is not performant, but is kept simple for illustrative purposes.
Defining interface methods#
As it currently stands, this GEMM kernel implementation cannot be used via the CUTLASS Operator API because it does not implement interface methods. Specifically, operators within the CUTLASS Operator API must inherit from and implement the cutlass.operators.Operator abstract class. This class has methods needed for many common operations performed when compiling and executing DSL kernels.
Certain providers (i.e., DSLs), such as CuTe DSL, provide an additional layer atop the Operator class to add convenience utilities for Operators using that provider. For example, the CuTe DSL provider in CUTLASS Operator API defines ops.providers.cutedsl.operator.CuteDslOperator, which adds utilities surrounding cute.compile() to add compile-time arguments needed for using TVM-FFI when it is enabled.
We will next walk through the steps in defining interface methods for this implementation
[ ]:
import itertools
import cutlass.operators as ops
We begin by defining a class to represent the Operator’s interface. As mentioned above, since this is a CuTe DSL Operator, our interface must inherit from and implement ops.providers.cutedsl.operator.CuteDslOperator.
The class must additionally be registered with the CuTe DSL provider via the @CuTeDSLProvider.register decorator so that the class can be considered when discovering Operators.
Each Operator subclass must define two class-level attributes:
supported_args_type: TheRuntimeArgumentssubclass this Operator supports (e.g.,GemmArguments). This allows providers to efficiently filter Operator classes before callinggenerate_operators().designed_for_min_cc: Minimum compute capability required by this Operator class. This indicates a broad minimum that this Operator implementation as a whole is designed for. Individual Operator instances present more finegrained information about which arch they support, underOperatorMetadata.supported_targets.
[3]:
@ops.providers.cutedsl.CuTeDSLProvider.register
class F64GemmOperator(ops.providers.cutedsl.operator.CuteDslOperator):
# Declare which RuntimeArguments subclass this Operator supports.
supported_args_type = ops.GemmArguments
# Minimum compute capability required by this Operator class.
designed_for_min_cc = 80
# Empty versions of interface methods. These will be implemented later, interspersed
# with notebook markdown. Normally, one would define them inline with the class definition.
def __init__(self, metadata: ops.OperatorMetadata):
pass
def _run(
self,
args: ops.GemmArguments,
artifact: ops.artifact.CompiledArtifact,
stream,
workspace=None,
):
pass
def _compile(
self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None
) -> ops.artifact.CompiledArtifact:
pass
@classmethod
def _generate_operators(
cls, metadata_filter, epilogue_args=None, target_sm=None, args=None
) -> list["F64GemmOperator"]:
pass
def _supports(
self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None
) -> ops.Status:
pass
def _get_workspace_size(
self, args: ops.GemmArguments
) -> ops.workspace.AllocationRequirement:
pass
The __init__ method of the class takes in a OperatorMetadata object from which it extracts the cta_tile_shape_mn. This is used to construct the Operator implementation object. We will discuss later how the OperatorMetadata object passed in here is constructed:
[4]:
def __init__(self, metadata: ops.OperatorMetadata):
# Using Python-2-style super() because we're defining this method outside of the class definition.
super(F64GemmOperator, self).__init__(metadata)
cta_tile_shape_mn = metadata.design.tile_shape[:2]
self.impl = F64GemmKernelImplementation(cta_tile_shape_mn)
Defining interfaces for compilation and execution#
The interfaces needed for compilation and execution are simple.
The compile method simply constructs a placeholder stream object and passes that and relevant arguments to self.cute_compile. This is a utility defined in the CuteDSLOperator abstract class that passes in compilation flags needed for certain options to cute.compile (e.g., TVM-FFI). The result is wrapped as a CompiledArtifact.
[5]:
def _compile(
self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None
) -> ops.artifact.CompiledArtifact:
stream = cutlass.cute.runtime.make_fake_stream()
return self.cute_compile(
self.impl,
args.A.tensor,
args.B.tensor,
args.out.tensor,
stream,
target_sm=target_sm,
)
Users define the _run method rather than the top-level run method (no leading underscore) that is used in interacting with operators. _run (1) extracts from args the arguments needed to run the JIT function, and (2) calls the JIT function passed in via artifact with these arguments.
[6]:
def _run(
self,
args: ops.GemmArguments,
artifact: ops.artifact.CompiledArtifact,
stream,
workspace=None,
):
stream = ops.utils.device.to_cuda_stream(stream)
compiled_gemm = artifact.compiled_obj
self.cute_run(compiled_gemm, args.A.tensor, args.B.tensor, args.out.tensor, stream)
Finally, since this Operator does not require any device workspace, we give it a simple get_workspace_size method that always returns an empty workspace. Note that this is the same as the base implementation of this method, and so it’s not strictly required in this case. We show it here for demonstration of the required methods if workspace is used.
[7]:
def _get_workspace_size(
self, args: ops.GemmArguments
) -> ops.workspace.AllocationRequirement:
return ops.workspace.AllocationRequirement.empty()
Defining interfaces for Operator generation#
We have implemented the interfaces needed for constructing the Operator interface, compiling it, and running it. We now must implement methods for generating the possible configurations of this Operator that the Operator class itself supports. This will be used in Operator discovery (e.g., via ops.get_operators()).
To do so, we write the generate_operators method. This takes in a binary function metadata_filter, epilogue arguments epilogue_args, and a compute capability cc. It returns a list of all instances of the Operator interface that support the epilogue_args, are compatible with the given cc, and which pass the metadata_filter.
The Operator class is responsible for defining what valid possible configurations (instances) of it can exist. In this example, the valid configurations involve a cross-product of row/column-major strides and two preset tile shapes. We create a nested loop over these options and create a OperatorMetadata corresponding to each unique configuration.
The generate_operators method must additionally filter the generated operators by passing it through a metadata_filter. This is a user-provided custom filter to filter generated metadata combinations. More information on metadata_filter is provided in other examples.
[8]:
@classmethod
def _generate_operators(
cls,
metadata_filter: Callable[[ops.OperatorMetadata], bool],
epilogue_args: ops.EpilogueArguments = None,
target_sm: ops.TargetSm | None = None,
args: ops.GemmArguments | None = None,
) -> list["F64GemmOperator"]:
# The tile shapes this Operator supports/exposes
supported_tile_shapes = [(32, 32, 1), (16, 16, 1)]
if epilogue_args is not None:
return []
row_major_stride = (0, 0, 1)
col_major_stride = (0, 1, 0)
stride_combos = list(
itertools.product([row_major_stride, col_major_stride], repeat=3)
)
# If args are provided, use the strides from the arguments
if args is not None:
if any(x.dtype != cutlass.Float64 for x in [args.A, args.B, args.out]):
return []
if any(not isinstance(x, ops.DenseTensor) for x in [args.A, args.B, args.out]):
return []
A_stride = ops.utils.tensor.normalized_major_stride(
args.A.shape, args.A.stride, prepend_zeros_to_rank=3
)
B_stride = ops.utils.tensor.normalized_major_stride(
args.B.shape, args.B.stride, prepend_zeros_to_rank=3
)
out_stride = ops.utils.tensor.normalized_major_stride(
args.out.shape, args.out.stride, prepend_zeros_to_rank=3
)
combo = (A_stride, B_stride, out_stride)
if combo not in stride_combos:
return []
# Override stride_combos
stride_combos = [combo]
divisibility = 1
def stride_name(stride):
return "T" if stride == row_major_stride else "N"
operators = []
for tile_shape in supported_tile_shapes:
for stride_A, stride_B, stride_out in stride_combos:
# Create OperandConstraints for A, B, and out tensors
a_attrs = ops.metadata.DenseTensorConstraints(
cutlass.Float64, stride_A, divisibility
)
b_attrs = ops.metadata.DenseTensorConstraints(
cutlass.Float64, stride_B, divisibility
)
out_attrs = ops.metadata.DenseTensorConstraints(
cutlass.Float64, stride_out, divisibility
)
layout_str = ops.utils.tensor.strides_to_layout_string(
stride_A, stride_B, stride_out
)
design = ops.metadata.BLASDesignMetadata(
mma_instruction_type=ops.mma.AmpereMma,
tile_shape=tile_shape,
cluster_shape=(1, 1, 1),
)
operands = ops.metadata.GemmOperandsMetadata(
A=a_attrs, B=b_attrs, out=out_attrs, accumulator_type=cutlass.Float64
)
name = f"F64GemmOperator_tile{tile_shape[0]}x{tile_shape[1]}_{layout_str}"
metadata = ops.OperatorMetadata(
operator_name=name,
operator_class=F64GemmOperator,
operands=operands,
design=design,
supported_targets=ops.TargetSm.get_supported_targets(design, operands),
)
if metadata_filter(metadata):
operators.append(cls(metadata))
return operators
We also add a method for indicating whether a Operator instance in question supports a set of arguments. The top-level Operator.supports method will already verify that the args and target_sm match the metadata with which this Operator instance was constructed. Here, we can define additional checks specific to this Operator. For instance, here we require that all operands to be of rank 3:
[9]:
def _supports(
self, args: ops.GemmArguments, target_sm: ops.TargetSm | None = None
) -> ops.Status:
if not (
len(args.A.shape) == 3 # A should be (L, M, K)
and len(args.B.shape) == 3 # B should be (L, K, N)
and len(args.out.shape) == 3 # out should be (L, M, N)
):
return ops.Status.fail("All operands must be rank 3.")
return ops.Status.success()
[10]:
# Assign methods to the class because we interspersed notebook markdown
# with the class definition. This is not needed in a real implementation.
F64GemmOperator.__init__ = __init__
F64GemmOperator._compile = _compile
F64GemmOperator._run = _run
F64GemmOperator._supports = _supports
F64GemmOperator._generate_operators = _generate_operators
F64GemmOperator._get_workspace_size = _get_workspace_size
Discovering instances of the Operator and using them#
The CUTLASS Operator API is now prepared to discover instances of this Operator interface just as was done in previous examples.
We add a small modification of using a metadata_filter to ensure that all returned operators are instances of the F64GemmOperator class we just implemented. This is needed only for example/testing purposes.
[11]:
import torch
torch.manual_seed(2025)
L, M, N, K = 1, 256, 1024, 128
A = torch.randint(-2, 3, (L, M, K), device="cuda", dtype=torch.float64)
B = torch.randint(-2, 3, (L, K, N), device="cuda", dtype=torch.float64)
out = torch.empty(L, M, N, device="cuda", dtype=torch.float64)
args = ops.GemmArguments(A, B, out, accumulator_type=torch.float64)
def is_f64gemm_operator(metadata):
return metadata.operator_class is F64GemmOperator
operators = ops.get_operators(args, metadata_filter=is_f64gemm_operator)
We can print off the names of the first few operators to see that they come from our recently-added Operator.
[12]:
print(operators[0].metadata.operator_name)
print(operators[1].metadata.operator_name)
F64GemmOperator_tile32x32_ttt
F64GemmOperator_tile16x16_ttt
We can evaluate and test the correctness of an instance of our Operator:
[13]:
operators[0].run(args)
torch.testing.assert_close(out, A @ B)
We can also test the limits of our Operator’s design space by providing a metadata filter that expects a CTA tile size M of 256, which is not exposed in the generate_operators method of our recently-added Operator. We expect no operators of type F64GemmOperator to be returned.
[14]:
def my_filter(metadata):
return (
is_f64gemm_operator(metadata)
and isinstance(metadata.design, ops.metadata.BLASDesignMetadata)
and metadata.design.tile_shape[0] == 256
)
operators_ctam256 = ops.get_operators(args, metadata_filter=my_filter)
# No operators should be found
assert len(operators_ctam256) == 0
A note on contributing operators to directory structure#
This example showed how to define an Operator inline and add it to the API for example purposes. This Operator doesn’t necessarily need to live within the API’s source code.
We welcome contributions of operators that do live within the CUTLASS API’s repository as well.
Kernels in the repository are organized based on the “provider” in which they are authored (i.e., the DSL). All operators corresponding to a given provider live a directory corresponding to that provider under cutlass/operators/providers. For example, CuTe DSL operators live under cutlass/operators/providers/cutedsl.
Each provider can organize operators differently. For CuTe DSL, operators are further split based on their logical operation, with GEMM operators under the cutlass/operators/providers/cutedsl/gemm directory.
We recommend separating the implementation of the Operator from its interface not just by using separate classes, as done in this example, but also by separating the implementation and interface into separate files. This makes it easier to update each without affecting the other.
For example, CuTe DSL GEMM operators have the following organization:
cutlass/operators/
providers/
cutedsl/
gemm/
sm100_persistent.py
implementations/
sm100_persistent_impl.py