Arguments and Operands#

An operation is expressed as a RuntimeArguments object that describes both the operation kind (e.g. GEMM) and the operands it runs on. Each operation has its own RuntimeArguments subclass. Operands are described by Operand subclasses such as DenseTensor and ScaledOperand. Tensor- and numeric-typed fields within Arguments and Operands accept any object matching the TensorLike / NumericLike protocols (e.g. torch.Tensor, cute.Tensor, cutlass.Numeric, torch.dtype).

RuntimeArguments#

class cutlass.operators.arguments.base.RuntimeArguments( *, performance: PerformanceControls | None = None, )#

Bases: object

Describes the operands and all other arguments passed to an Operator at runtime.

It contains runtime operands (usually tensors) passed to the operation, as well as any custom epilogue fusions, and runtime performance controls.

This is an abstract base class, whose subclass describes the operation type itself (e.g. GemmArguments, GroupedGemmArguments, etc.). All operators that implement the same operation type accept the same RuntimeArguments subclass.

performance: PerformanceControls | None = None#: Optional runtime performance controls passed to the Operator

class cutlass.operators.arguments.base.PerformanceControls#

Bases: object

Optional runtime performance controls passed to the Operator.

Some operators may support performance options that can be controlled at runtime. This class is the general container for all such controls.

class cutlass.operators.arguments.gemm.GemmProblemSize(M: int, N: int, K: int, L: int)#

Bases: NamedTuple

Problem size for a GEMM operation.

A GEMM with problem size (M, N, K, L) has operands with the following shapes (batch, rows, columns):

A: (L, M, K)
B: (L, K, N)
out: (L, M, N)

M: int#: Number of rows in A and out

N: int#: Number of columns in B and out

K: int#: Number of columns in A and rows in B

L: int#: Number of batches of matrix multiplications

class cutlass.operators.arguments.gemm.GemmArguments( A: TensorLike | Operand, B: TensorLike | Operand, out: TensorLike | Operand, accumulator_type: NumericLike, epilogue: EpilogueArguments | None = None, )#

Bases: RuntimeArguments

Arguments for a Generalized Matrix Multiplication (GEMM) operation: out = A @ B.

The tensors must be all rank-3 or all rank-2. * L: Number of batches * M: Number of rows in A and out * K: Number of columns in A and rows in B * N: Number of columns in B and out

For convenience, construction for A, B, and out that are operands to a dense GEMM can be passed in without wrapping them in DenseTensor.

GemmArguments(A, B, out, accumulator_type)
# is equivalent to:
GemmArguments(DenseTensor(A), DenseTensor(B), DenseTensor(out), accumulator_type)

Other operand types must explicitly wrap tensors in a Operand subclass. For example, a scaled GEMM can be constructed as:

GemmArguments(
    ScaledOperand(A, ScaleATensor, scale_mode, scale_swizzle),
    ScaledOperand(B, ScaleBTensor, scale_mode, scale_swizzle),
    out, # No need to wrap in a `DenseTensor`
    accumulator_type,
)

A: Operand#: Input tensor A of shape (L, M, K) or (M, K)

B: Operand#: Input tensor B of shape (L, K, N) or (K, N)

out: Operand#: Output tensor C of shape (L, M, N) or (M, N)

accumulator_type: NumericLike#: Data type of the accumulator

epilogue: EpilogueArguments | None#: Optional custom epilogue fusion to be performed after the GEMM

property problem_size: GemmProblemSize#: Problem size for a GEMM operation.

class cutlass.operators.arguments.grouped_gemm.GroupedGemmArguments( A: TensorLike | Operand, B: TensorLike | Operand, out: TensorLike | Operand, accumulator_type: NumericLike, offsets: TensorLike | Operand, epilogue: EpilogueArguments | None = None, )#

Bases: RuntimeArguments

Arguments for a grouped GEMM operation.

A grouped GEMM performs a series of independent GEMM operations, where each GEMM can have different matrix dimensions. The most basic formulation of a grouped GEMM is one that takes lists of tensors for each of A, B, and out, and computes the following:

for i in range(problems_in_group):
    out[i] = A[i] @ B[i]

While the abstract formulation uses lists of tensors, in practice, the tensors are often packed together contiguously in memory in a single tensor for each of A, B, and out. In this case, an offsets tensor delineates the ending positions of each problem in the group within those contiguous tensors.

Currently-supported variants of a grouped GEMM are:

Contiguous offset 2D-3D grouped GEMM:

A is a tensor of shape (TotalM, K) or (1, TotalM, K)
B is a tensor of shape (problems_in_group, K, N)
out is a tensor of shape (TotalM, N) or (1, TotalM, N)
offsets is a tensor delineating the ending positions of each problem in the group.

start = 0
for i in range(problems_in_group):
    end = offsets[i]
    out[start:end, :] = A[start:end, :] @ B[i, :, :]
    start = end

A: Operand#: Input tensor A

B: Operand#: Input tensor B

out: Operand#: Output tensor

accumulator_type: NumericLike#: Data type of the accumulator

offsets: Operand#: Offsets tensor delineating the ending positions of each problem in the group

epilogue: EpilogueArguments | None#: Optional custom epilogue fusion to be performed after the GEMM

class cutlass.operators.arguments.epilogue.EpilogueArguments(epilogue_fn: Callable | str, **kwargs)#

Bases: object

Describes a user-defined epilogue that is fused on top of the operation described by the primary RuntimeArguments.

An epilogue fusion is a custom function that performs tensor-level transformations on the result of a matrix multiplication. This transformation is fused into the kernel’s epilogue, which stores the final output.

EpilogueArguments encapsulates the epilogue function describing the transformation along with its arguments.

To support flexible definition of epilogues, 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. kwargs must contain Tensors or scalars for all operands and outputs in the provided epilogue.

Structure of ``epilogue_fn``

The epilogue_fn is a function describing the custom transformation on the accumulator tensor (intermediate result before the epilogue).

The general structure of these functions is:

def custom_epi_name(accum, *args) -> TensorType | tuple[TensorType, ...]:
    '''Compute the epilogue.

    # Args:
        accum (TensorType): Result of the primary operation (e.g.
                        ``A @ B`` for a GEMM) before the epilogue.
        *args: Additional tensors or scalars used in the epilogue
                (e.g. aux tensors).

    # Returns:
        At least one tensor resulting from the epilogue computation.
    '''
    # Do some compute
    return D  # and potentially other values

epilogue_fn must be a Python callable (or its string representation) that must satisfy the following constraints:

Takes a first positional argument named accum – the result of the operation just before the epilogue. For a GEMM, accum = A @ B.
Returns at least one tensor resulting from the epilogue. Currently the return list must contain at least one output named D.
Each argument following accum is a tensor or scalar to be loaded.
Each variable in the return statement is a tensor or scalar to be stored.
Operations are represented in static single assignment (SSA) form. This means that each variable can be assigned exactly once.

The underlying implementation of the epilogue in the kernel will determine how operands are loaded and stored.

Structure of ``kwargs``

kwargs must contain sample 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=...
)

property parameters: list[cute.Tensor | Numeric]#: Returns the list of input and output parameters of the epilogue.

property parameter_names: list[str]#: Returns the list of names of the input and output parameters of the epilogue.

to_tensor_wrappers(permute: list[int] | None = None)#: Converts the input and output parameters of the epilogue to TensorWrappers.

trace( accumulator_shape: tuple[int, ...], accumulator_type: Numeric, )#

Traces the epilogue function and generates an internal representation of the epilogue.

Parameters:

accumulator_shape (tuple[int, ...]) – The shape of the accumulator tensor. For example, for a GEMM, this would be the shape of the output tensor.
accumulator_type (Numeric) – The datatype of the accumulator tensor.

Operands#

class cutlass.operators.arguments.base.Operand#

Bases: ABC

Base class for all operands to Operators, which encapsulates one or more TensorLike objects.

In the most basic case, an Operand enacapsulates a single tensor.

In more complex cases, an Operand may encapsulate multiple tensors that encapsulate a single logical operand. For instance, a ScaledOperand encapsulates a quantized and scale tensor, that together reconstruct the logical value of the operand.

final copy() → Operand#: Returns a copy of the operand. Does not copy the underlying tensor.

class cutlass.operators.arguments.operand.DenseTensor(tensor: cutlass.operators.typing.TensorLike)#

Bases: Operand

An operand encapsulating a simple, single dense tensor.

class cutlass.operators.arguments.operand.ScaledOperand( quantized: TensorLike | DenseTensor, scale: TensorLike | DenseTensor, mode: ScaleMode | tuple[int, ...], swizzle: ScaleSwizzleMode, )#

Bases: Operand

An operand whose logical value is a quantized tensor multiplied by a tensor of scale factors.

A ScaledOperand encapsulates physical quantized and scale tensors that together represent the logical value scale * quantized: each entry of scale broadcasts and multiplies a contiguous block of quantized, where the block shape is given by mode.

This Operand is typically used to express narrow-precision formats (such as OCP MXFP8 / MXFP4 and NVIDIA NVFP4), which store the data as a quantized narrow-precision tensor multiplied by a separate tensor of scale factors that recover dynamic range.

The scale mode and swizzle describe the physical layout of the scale tensor.

The mode describes how scale factors are broadcast over the quantized tensor, and therefore determines its size. It’s usually dictated by the data format you seek to represent – MXFP8 / MXFP4 / NVFP4 require specific scale modes. See ScaleMode for more details.

The swizzle describes the in-memory layout of the scale tensor. It’s usually dictated by the hardware architecture – for example, Blackwell block-scaled MMAs require a specific swizzle layout, and most kernels using them require the input to be pre-arranged in this layout. See ScaleSwizzleMode for more details.

Example

quantized_A = torch.randn(M, K, dtype=torch.float8_e4m3fn, device="cuda")
scale_A = torch.randn(M, K // 32, dtype=torch.float8_e8m0fnu, device="cuda")
A = ScaledOperand(
    quantized_A, scale_A, ScaleMode.Blockwise1x32, ScaleSwizzleMode.Swizzle32x4x4,
)

Type markers#

Type markers for CUTLASS Operator API field annotations.

cutlass.operators.typing.TensorLike#

Tensor-like inputs accepted by the Operator API.

Includes:

Any DLPack-compatible tensor (torch.Tensor, jax.Array, numpy.ndarray, …).
cutlass.cute.Tensor (CuTe DSL host tensor; does not implement DLPack but is handled natively by TensorWrapper).
cutlass.operators.utils.tensor.TensorWrapper (implements DLPack).

alias of _SupportsDLPack | Tensor | TensorWrapper

class cutlass.operators.typing.NumericLike(*args, **kwargs)#

Bases: Protocol

Type marker for fields that accept numeric-like types.

Fields annotated with NumericLike accept the following:

cutlass.Numeric
torch.dtype