Using fake tensors with the CUTLASS Operator API

Fake tensors (e.g., torch's FakeTensor) are useful for describing the properties of a tensor without actually allocating backing data.

This example shows how fake tensors can be used within the CUTLASS Operator API for discovering and compiling a GEMM Operator.

python

import torch

import cutlass.operators as ops

torch.manual_seed(2025)

if not (status := ops.utils.device.device_or_env_supports("80")):
    print(
        f"This notebook requires a GPU with compute capability >= 80.\n{status.error}"
    )
    import sys

    sys.exit(0)

We first set up operands A, B, and out in torch's FakeTensorMode. These will have all the properties needed for CUTLASS Operator API to construct the internal representations of tensors used for discovering and compiling operators.

python

M, N, K = 128, 256, 512

with torch._subclasses.fake_tensor.FakeTensorMode():
    A_fake = torch.randint(-2, 3, (M, K), device="cuda", dtype=torch.float16)
    B_fake = torch.randint(-2, 3, (K, N), device="cuda", dtype=torch.float16)
    out_fake = torch.empty(M, N, device="cuda", dtype=torch.float16)

print(A_fake)
print(B_fake)
print(out_fake)

We can now use these fake tensors to create GemmArguments, and use these to discover and compile a compatible operator. Note that the same APIs are used in creating GemmArguments as would be used if using "real" tensors.

python

args_fake = ops.GemmArguments(A_fake, B_fake, out_fake, accumulator_type=torch.float32)

target_sm = ops.utils.device.device_or_env_target_sm()
operators = ops.get_operators(args_fake, target_sm=target_sm)
assert len(operators) > 0

operator = operators[0]
compiled_artifact = operator.compile(args_fake)

The operator and compiled_artifact discovered using fake tensors above can now used for running the operator using real tensors.

python

# Create real tensors
A_real = torch.randint(-2, 3, (M, K), device="cuda", dtype=torch.float16)
B_real = torch.randint(-2, 3, (K, N), device="cuda", dtype=torch.float16)
out_real = torch.empty(M, N, device="cuda", dtype=torch.float16)

args_real = ops.GemmArguments(A_real, B_real, out_real, accumulator_type=torch.float32)

# Run the operator using the compiled_artifact from resulting
# from compiling with fake tensors.
operator.run(args_real, compiled_artifact)

python

ref = A_real @ B_real
torch.testing.assert_close(out_real, ref)