Note
This page is a rendered version of 001_gemm_with_fused_epilogue.ipynb available on GitHub.
Custom epilogue fusions for GEMMs#
[1]:
import cutlass.operators as ops
if not (status := ops.utils.device.device_or_env_supports("100f")):
print(f"This notebook requires a Blackwell GPU (sm_100f family).\n{status.error}")
import sys
sys.exit(0)
CUTLASS Operator API provides flexible epilogue fusion support – users can specify high-level tensor transformations that can be fused into an existing Operator. This is useful for fusing activations, storing auxiliary tensors for backprop, elementwise ops, broadcasts, etc.
For those familiar with the legacy CUTLASS Python API’s epilogue visitor tree frontend, much of the interface is shared.
CUTLASS Operator API enables one to express an epilogue using a function operating at the torch.Tensor-level, and has tooling to automatically add this to operators supporting the provided function.
For example, in PyTorch one might write the following to compute a GEMM + epilogue:
[2]:
import torch
torch.manual_seed(2025)
L, M, N, K = 1, 1024, 1024, 1024
A = torch.randint(-2, 3, (L, M, K), device="cuda", dtype=torch.float16)
B = torch.randint(-2, 3, (L, K, N), device="cuda", dtype=torch.float16)
C = torch.randint(-2, 3, (L, M, N), device="cuda", dtype=torch.float16)
def my_epilogue(accum, C, alpha, beta, extra_scalar):
Aux = (alpha * accum) + (beta * C)
D = extra_scalar * Aux
return D, Aux
alpha, beta, extra_scalar = 1.0, 2.0, 0.5
D, Aux = my_epilogue(A @ B, C, alpha, beta, extra_scalar)
Now suppose you want to compose this my_epilogue computation with a CUTLASS Python kernel that performs a regular GEMM. CUTLASS Operator API allows doing so, where supporting GEMM Operators can fuse the computation into the epilogue before storing the final output.
To do so, one defines EpilogueArguments consisting of the epilogue function to compute (or a string representation of it) along with arguments corresponding to each input and output of the function (except for accum):
[3]:
# Allocate buffers for D and Aux
D_, Aux_ = [
torch.empty((L, M, N), device="cuda", dtype=torch.float16) for _ in range(2)
]
epi_args = ops.EpilogueArguments(
my_epilogue,
C=C,
alpha=alpha,
beta=beta,
extra_scalar=extra_scalar,
D=D_,
Aux=Aux_,
)
These arguments can be added to GemmArguments and passed in to get_operators() for use when retrieving compatible operators:
[4]:
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
target_sm = ops.utils.device.device_or_env_target_sm()
operators = ops.get_operators(args, target_sm=target_sm)
assert len(operators) > 0
Each of the operators returned by get_operators can be compiled and executed just the same with these new arguments, as it was in examples without epilogue fusion. For example, using the first operator:
[5]:
operators[0].run(args)
torch.testing.assert_close(D, D_)
torch.testing.assert_close(Aux, Aux_)
How the epilogue fusion API works#
To support specifying an epilogue via a Python function, an Operator needs some mechanism to:
Detect the operations in the epilogue function
Determine if the Operator can support the operations
Emit code to perform these operations within the kernel
Step 1 listed above does not depend on the Operator and its implementation (e.g., DSL), while steps 2 and 3 depend on the Operator and/or its implementation.
CUTLASS Operator API separates these components so that step 1 takes place at the API level and steps 2 and 3 take place in the Operator. This process is visualized below. We will walk through each step in greater detail.
+------------------------------------+
| def epi(accum, alpha, beta, C): |
| D = (accum * alpha) + (beta * C) | 1. Define epilogue via a Python function
| return D |
+------------------------------------+
|
|
|
GemmArguments(..., 2. Pass epilogue function, operands, and outputs
epilogue=EpilogueArguments( to EpilogueArguments constructor,
epi, alpha=alpha, beta=beta, C=C)) and add this to the GemmArguments. Under the
| hood, this parses the Python AST of the
| epilogue function to produce a DAG of load,
| store, and compute nodes.
V
+-----------------------------------------+
| Intermediate DAG representation |
| =============================== |
| |
| Store() |
| | |
| Add() |
| / \ |
| / \ |
| / \ |
| Mul() Mul() |
| / \ / \ |
| AccFetch() | Load(C) \ |
| | \ |
| Load(alpha) Load(beta) |
| |
+-----------------------------------------+
/ | \
/ | \ 3. Individual Operator classes use the DAG representation
/ | \ to determine if the Operator class supports the DAG.
Operator 0 Operator 1 Operator 2 If so, the Operator class emits DSL-level operations
epilogue epilogue epilogue needed to compute the epilogue DAG alongside the
emitter emitter emitter basic operation of the kernel (e.g., GEMM).
| | |
| | |
V V V
Defining an epilogue via a Python function#
Epilogue fusion patterns are defined by users in Python functions that perform Tensor-level transformations – using a torch.Tensor (for example) – on the resul of a matrix multiplication.
The structure of these functions is as follows:
def custom_epi_name(accum, *args) -> Union[TensorType, tuple[TensorType]]:
"""
:param accum: result of matrix multiplication, convolution, etc. before the epilogue
:type accum: TensorType
:param args: additional arguments to be used in the epilogue (e.g., aux tensors)
:type args: list[Union[TensorType, ScalarType]]
:returns: at least one tensor resulting from the operation of the epilogue
:rtype: Union[TensorType, tuple[TensorType]]
"""
# Do some compute
return D # and potentially other values
The user defines a custom epilogue via a Python function that must do at least the following:
Take in a first positional argument named
accumthat represents the result of operation just before the epilogue is to be performed. For example, in a GEMM,accum = A @ B.Return at least one tensor that results from computing the epilogue. Currently, the return list must contain at least one output named
D, though this constraint may be loosened in the future.
Each additional argument following accum in the function definition is expected to be either a Tensor or scalar to be loaded. Each variable in the return statement represents a Tensor or scalar to be stored. The underlying implementation of the epilogue in the kernel will determine how operands are loaded and stored.
Compute operations are represented in static single assignment (SSA) form. This means that each variable can be assigned exactly once. Operations currently supported are:
Tensor-tensor elementwise addition, subtraction, multiplication, and division
Scalar broadcasts via addition, subtraction, multiplication, and division
Row broadcasts: load a bias of shape
(N,)or(1, N)and broadcast across MColumn broadcasts: load a bias of shape
(M,)or(M, 1)and broadcast across NCombined row and column broadcasts in the same epilogue
Binary elementwise minimum and maximum
Predefined elementwise activation functions (e.g., ReLU, sigmoid, tanh)
Operations that are not yet supported include:
Reductions (row, column, scalar)
If attempting to use unsupported operations, no operators will be found in the call to get_operators.
Violations to SSA or use of unexpected operators will be flagged with an exception when parsing the AST of the custom epilogue.
Examples of epilogues fitting these patterns are given below. We will show full, runnable examples at the end of this notebook.
def relu_aux_store(accum, alpha, C):
# Note that the function definition itself does not indicate the types and
# ranks of alpha and C. Thus, one cannot tell whether the epilogue is performing
# broadcasts or elementwise operations until actual arguments or metadata are
# provided to the epilogue. See below for details.
F = (accum * alpha) + (C * 2.0) # Constant beta of 2.0
D = relu(F)
return D, F
def aux_normalize(accum, aux):
D = accum / aux
return D
def bias_add(accum, row_bias, col_bias):
# row_bias has shape (N,) and col_bias has shape (M, 1).
# The broadcast type is determined automatically from the tensor shapes
# provided to EpilogueArguments.
D = accum + row_bias + col_bias
return D
Additional information about each operand and output must be provided by the user when constructing EpilogueArguments, as we will discuss below. This additional information is necessary for fully defining the operations being performed – without knowledge of whether alpha is a scalar or a Tensor, we cannot determine whether multiplication by alpha is a broadcasted or elementwise operation.
Constructing epilogue arguments#
EpilogueArguments encapsulate the arguments needed to determine the functional operation of a fused epilogue.
A user must provide in the construction of EpilogueArguments tensors for all operands and outputs of the epilogue. However, unlike arguments for basic operations (e.g., GEMM), the full set of operands needed to be specified for an epilogue pattern depends upon the custom epilogue defined by the user.
Therefore, EpilogueArguments is defined generically as taking in an epilogue_fn and additional kwargs. Under the hood, the AST for epilogue_fn is parsed to determine the operands and outputs of the epilogue. The user is required to provide in kwargs Tensors or scalars for all operands and outputs in the provided epilogue.
For example, with an epilogue of:
def my_epi(accum, alpha, C, beta):
F = (accum * alpha) + (C * beta)
D = relu(F)
return D, F
A user would need to construct epilogue arguments as follows:
epi_args = EpilogueArguments(my_epi, alpha=..., C=..., beta=..., D=..., F=...)
After verifying that all required operands and outputs are present, the constructor to EpilogueArguments will perform additional passes on the AST of epilogue_fn using the provided inputs to generate an internal DAG representing the epilogue. This DAG structure is attached to EpilogueArguments for use as they are passed through a call to get_operators.
Discovering operators that support the epilogue pattern#
The call to get_operators(args) will return any operators that support the provided GemmArguments. Since the GemmArguments constructed above now include EpilogueArguments, returned operators must support the provided epilogue.
Under the hood of get_operators(), each Operator class will determine in its generate_operators() method whether it supports the provided EpilogueArguments. It can do so by traversing the DAG that resulted from the construction of EpilogueArguments to find the operations that compose the epilogue. Assuming that the Operator can support the DAG, it must then add to the source for the kernel any operations needed to support the DAG. An example of how this is done
generically for an SM100 CuTe DSL GEMM is provided in `cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py <../cutlass/operators/providers/cutedsl/gemm/sm100_static_persistent_efc.py>`__.
Example epilogues#
We now provide various examples of adding custom epilogues to GEMM operators targeting SM100. A broader set of epilogue examples are available in `test/integration/test_gemm_epilogue_fusion.py <../test/integration/test_gemm_epilogue_fusion.py>`__.
Auxiliary input and output tensors#
[6]:
from cutlass.operators.fusion.activation import relu
def relu_aux_store(accum, alpha, C):
F = (accum * alpha) + (C * 2.0) # Constant beta
D = relu(F)
return D, F
C = torch.randint(-2, 3, (L, M, N), device="cuda", dtype=torch.float16)
alpha = 3.0
D = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
F = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
epi_args = ops.EpilogueArguments(relu_aux_store, alpha=alpha, C=C, D=D, F=F)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
operators = ops.get_operators(args, target_sm=target_sm)
assert len(operators) > 0
operators[0].run(args)
D_ref, F_ref = relu_aux_store(A @ B, alpha, C)
torch.testing.assert_close(D, D_ref)
torch.testing.assert_close(F, F_ref)
Keyword functions and returning accumulator#
[7]:
def relu_scale_return_acc(accum, alpha, beta, C, scale):
F = relu((accum * alpha) + (C * beta))
D = F * scale
return D, F, accum
C = torch.randint(-2, 3, (L, M, N), device="cuda", dtype=torch.float16)
alpha = 1.0
beta = 2.0
scale = 0.5
D = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
F = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
accum = torch.empty((L, M, N), device="cuda", dtype=torch.float32)
epi_args = ops.EpilogueArguments(
relu_scale_return_acc,
alpha=alpha,
beta=beta,
C=C,
scale=scale,
D=D,
F=F,
accum=accum,
)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
operators = ops.get_operators(args, target_sm=target_sm)
assert len(operators) > 0
operators[0].run(args)
D_ref, F_ref, accum_ref = relu_scale_return_acc(A @ B, alpha, beta, C, scale)
torch.testing.assert_close(D, D_ref)
torch.testing.assert_close(F, F_ref)
torch.testing.assert_close(accum, accum_ref.to(accum.dtype))
Passing a string representation of the function#
EpilogueArguments can additionally be constructed using a string representation of the epilogue function:
[8]:
epi_str = "def epi(accum, alpha, beta, C): F = (accum * alpha) + (C * beta); D = relu(F); return D, F"
C = torch.randint(-2, 3, (L, M, N), device="cuda", dtype=torch.float16)
alpha = 1.0
beta = 0.5
D = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
F = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
epi_args = ops.EpilogueArguments(epi_str, alpha=alpha, beta=beta, C=C, D=D, F=F)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
operators = ops.get_operators(args, target_sm=target_sm)
assert len(operators) > 0
operators[0].run(args)
F_ref = (A @ B) * alpha + (C * beta)
D_ref = torch.relu(F_ref)
torch.testing.assert_close(D, D_ref)
torch.testing.assert_close(F, F_ref)
Row and column broadcasts#
A common epilogue pattern is adding a per-row or per-column bias to the GEMM result. The broadcast type is determined automatically from the shape of the tensor provided to EpilogueArguments:
Row broadcast: bias has shape
(N,)or(1, N)— broadcast across all M rows.Column broadcast: bias has shape
(M, 1)— broadcast across all N columns.
Both types can be combined in the same epilogue:
[9]:
def broadcast_epilogue(accum, row_bias, col_bias, alpha):
# row_bias shape (N,): same bias added to every row.
# col_bias shape (M, 1): same bias added to every column.
# The broadcast type is inferred from the shapes provided
# to EpilogueArguments.
D = (accum * alpha) + row_bias + col_bias
return D
# Row bias: one value per column (N).
row_bias = torch.randn(N, device="cuda", dtype=torch.float16)
# Column bias: one value per row (M), with an extra trailing dimension.
col_bias = torch.randn(M, 1, device="cuda", dtype=torch.float16)
alpha = 1.0
D = torch.empty((L, M, N), device="cuda", dtype=torch.float16)
epi_args = ops.EpilogueArguments(
broadcast_epilogue,
row_bias=row_bias,
col_bias=col_bias,
alpha=alpha,
D=D,
)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
operators = ops.get_operators(args, target_sm=target_sm)
assert len(operators) > 0
operators[0].run(args)
D_ref = broadcast_epilogue(A @ B, row_bias, col_bias, alpha)
torch.testing.assert_close(D, D_ref)
Failure examples#
The following are examples of constructing EpilogueArguments that are expected to fail.
[10]:
####################################################
# Epilogues must take in an accumulator
####################################################
def fail_missing_accum(alpha, beta, C):
D = C * beta
return D
try:
epi_args = ops.EpilogueArguments(
fail_missing_accum,
alpha=alpha,
beta=beta,
C=C,
D=D,
)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
except Exception as e:
# "accum must be an input to the epilogue function"
print(e)
accum must be an input to the epilogue function
[11]:
####################################################
# Epilogues must return an output named D
####################################################
def fail_missing_D(accum, alpha, beta, C):
F = (accum * alpha) + (C * beta)
return F
try:
epi_args = ops.EpilogueArguments(fail_missing_D, alpha=alpha, beta=beta, C=C, F=F)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
except Exception as e:
# "Output node D is not found in the epilogue function
print(e)
Output node D is not found in the epilogue function.
[12]:
####################################################
# Epilogues must use single-static assignment (SSA)
####################################################
def fail_ssa(accum):
tmp = accum * 2.0
# Redefine tmp, which violates SSA form.
tmp = tmp - 1.0
D = tmp / 4.0
return D, tmp
try:
epi_args = ops.EpilogueArguments(fail_ssa, D=D, tmp=F)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
except Exception as e:
# "Variable 'tmp' cannot be defined twice."
print(e)
Variable 'tmp' cannot be defined twice.
[13]:
####################################################
# Must provide all operands and outputs to
# EpilogueArguments
####################################################
def my_epi(accum, alpha, beta, C):
F = (accum * alpha) + (C * beta)
D = relu(F)
return D
try:
# Missing D
epi_args = ops.EpilogueArguments(my_epi, alpha=alpha, beta=beta, C=C)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
except Exception as e:
# "Argument D is not provided in the kwargs of the EpilogueArguments constructor"
print(e)
try:
# Missing alpha
epi_args = ops.EpilogueArguments(my_epi, beta=beta, C=C, D=D)
args = ops.GemmArguments(
A=A, B=B, out=D, accumulator_type=torch.float32, epilogue=epi_args
)
except Exception as e:
# "Argument alpha is not provided in the kwargs of the EpilogueArguments constructor"
print(e)
Argument D is not provided in the kwargs of the EpilogueArguments constructor
Argument alpha is not provided in the kwargs of the EpilogueArguments constructor