Best practices for reducing host-side latency

Overall performance depends on both device performance (i.e., that of the operator) and host performance (i.e., that of the runtime). This notebook focuses on the latter: techniques to minimize any overheads incurred from the CUTLASS API and underlying DSL runtimes.

This notebook does not discuss techniques for improving device-side performance. A future notebook may cover this topic.

Note: Latency measurements can vary from system to system. You may see different results on your system than shown in the pre-populated fields of this notebook.

python

import time

import torch

import cutlass.operators as ops

python

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 start with boilerplate initial setup to create tensors and pick an operator.

For the purposes of this notebook, we use a very small GEMM size of M=N=K=128 and L=1. This small size is chosen to magnify the impact of host latency on end-to-end performance so as to better illustrate the effect of the techniques described below.

python

warmup_iterations = 10
profiling_iterations = 100
total_iterations = warmup_iterations + profiling_iterations

# Use a small problem size to showcase host overheads
L, M, N, K = 1, 128, 128, 128

# We use different operands in each iteration. Though not particularly relevant for
# host latency, this is a best practice when benchmarking GPU kernels to avoid
# unrealistic caching effects.
As = [
    torch.randint(-1, 2, (M, K), device="cuda", dtype=torch.float16)
    for _ in range(total_iterations)
]
Bs = [
    torch.randint(-1, 2, (K, N), device="cuda", dtype=torch.float16)
    for _ in range(total_iterations)
]
outs = [
    torch.empty((M, N), device="cuda", dtype=torch.float16)
    for _ in range(total_iterations)
]

# Construct arguments outside of the benchmarking loop. We will later also consider
# cases in which they are constructed inside the benchmarking loop.
args = [
    ops.GemmArguments(A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float32)
    for i in range(total_iterations)
]

references = [(As[i] @ Bs[i]).to(outs[i].dtype) for i in range(total_iterations)]

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

operator = operators[0]

We next set up a basic benchmarking routine.

python

def benchmark(
    label, code, warmup_it=warmup_iterations, profiling_it=profiling_iterations
):
    total_it = warmup_it + profiling_it
    assert total_it <= total_iterations, (
        f"Benchmark-local iteration count must be less than or equal to total iterations: {total_it} > {total_iterations}"
    )
    # warmup
    rets = [None] * total_it
    for i in range(warmup_it):
        rets[i] = code(i)
    torch.cuda.synchronize()

    start = time.time()
    for i in range(profiling_it):
        idx = warmup_it + i
        rets[idx] = code(idx)
    torch.cuda.synchronize()
    end = time.time()

    avg_time = (end - start) / profiling_it
    print(f"[{label:<30}] avg of {profiling_it} iterations: {avg_time:1.3e} seconds")
    return avg_time, rets

We now describe techniques for reducing host latency:

Compile once, run many times
Bypassing checks for argument-Operator compatibility
Using CUDA Graphs
Using TVM FFI

These techniques are complementary and should be used together when applicable for an application.

Compile once, run many times

The operator.run method takes in an optional compiled_artifact argument of type ops.artifact.CompiledArtifact. When this argument is set, the operator will directly use the precompiled function within compiled_artifact. When it is not set, the call to operator.run will JIT compile the operator on each invocation.

Precompiling the operator is critical to achieving good performance.

python

stream = torch.cuda.current_stream()


def no_compiled_artifact(i: int):
    return operator.run(args[i], stream=stream)


# Compile the Operator once, reuse for each iterations
compiled_artifact = operator.compile(args[0])


def with_compiled_artifact(i: int):
    return operator.run(args[i], stream=stream, compiled_artifact=compiled_artifact)

python

time_no_artifact, _ = benchmark(
    "Without compiled artifact", no_compiled_artifact, warmup_it=2, profiling_it=5
)
time_w_artifact, _ = benchmark(
    "With compiled artifact", with_compiled_artifact, warmup_it=2, profiling_it=5
)

Bypassing checks for argument-Operator compatibility

By default, the call to operator.run will check if the Operator supports the provided arguments. Under the hood, this invokes operator.supports(args).

While these checks are helpful for catching incompatible arguments, they are performed in Python, and thus can add to host overhead.

When confident that arguments will be compatible with an Operator, one should bypass the supports check in operator.run by setting the optional assume_supported_args argument to True.

python

def with_supports_check(i: int):
    return operator.run(
        args[i],
        compiled_artifact=compiled_artifact,
        stream=stream,
        assume_supported_args=False,
    )


def without_supports_check(i: int):
    return operator.run(
        args[i],
        compiled_artifact=compiled_artifact,
        stream=stream,
        assume_supported_args=True,
    )

python

time_w_supports, _ = benchmark("With supports check", with_supports_check)
time_wo_supports, _ = benchmark("Bypass supports check", without_supports_check)
print(f"Speedup with skip supports: {time_w_supports / time_wo_supports:.2f}x")

CUDA Graphs

CUDA Graphs allow a sequence of GPU operations to be defined as a dependency graph and then launched as a single unit, significantly reducing CPU launch overhead and enabling whole-graph optimizations.

CUTLASS API supports CUDA Graphs usage with PyTorch as usual.

The operator compilation must happen outside the CUDA graph. Then, we create a graph using usual PyTorch idioms to launch a kernel several times on the graph's stream.

python

num_launches = 20

# Create a CUDA Graph to run our compiled Operator N times
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    ### NOTE! Operator compilation must happen outside the graph
    ### operator.compile(args)

    # Run N iterations of our compiled Operator on the current stream
    for i in range(num_launches):
        operator.run(
            args[i],
            compiled_artifact=compiled_artifact,
            stream=torch.cuda.current_stream(),
            assume_supported_args=True,
        )

This records/captures all the kernel launches to the CUDA Stream associated with the graph g, without actually launching them. Once captured, we can replay the graph.

Note that graph replay will only replay the kernel launches placed on the graph's stream

During graph capture, we must be careful to capture to the correct stream (torch.cuda.current_stream() under the graph context)
Any other preparatory work on the host and arguments passed in from Python are cached during the capture. Changing them would require re-capturing the graph

python

# Replay captured graph and check first result
g.replay()

torch.testing.assert_close(outs[0], references[0])

Let's compare the timing:

python

def without_cuda_graph(x: int):
    for i in range(num_launches):
        operator.run(
            args[i],
            compiled_artifact=compiled_artifact,
            stream=torch.cuda.current_stream(),
            assume_supported_args=True,
        )


def with_cuda_graph(x: int):
    g.replay()


time_wo_cuda_graph, _ = benchmark(
    f"{num_launches} launches without CUDA Graph",
    without_cuda_graph,
    warmup_it=0,
    profiling_it=1,
)
time_w_cuda_graph, _ = benchmark(
    f"{num_launches} launches with CUDA Graph",
    with_cuda_graph,
    warmup_it=0,
    profiling_it=1,
)

print(f"Speedup with CUDA Graph: {time_wo_cuda_graph / time_w_cuda_graph:.2f}x")

TVM FFI

Apache TVM FFI is an open ABI and FFI for machine learning systems. When available, CUTLASS API uses Apache TVM-FFI under the hood as its interface for invoking compiled DSL kernels from Python.

TVM FFI is enabled by default in CUTLASS API, and is recommended for best performance.

ops.GlobalOptions().use_tvm_ffi controls whether or not TVM-FFI will be used by CUTLASS API.

python

print(ops.GlobalOptions().use_tvm_ffi)

If for some reason you do not wish to use it, this section demonstrates how, you can set this to False. No other change is needed. The below code compares the performance with and without TVM-FFI.

python

original_use_tvm_ffi = ops.GlobalOptions().use_tvm_ffi

ops.GlobalOptions().use_tvm_ffi = True


def run_iteration(i):
    args = ops.GemmArguments(
        A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float16
    )
    return operator.run(
        args,
        compiled_artifact=compiled_artifact,
        stream=torch.cuda.current_stream(),
        assume_supported_args=True,
    )


def create_arguments(i: int):
    return ops.GemmArguments(
        A=As[i], B=Bs[i], out=outs[i], accumulator_type=torch.float16
    )


args_creation_on, args = benchmark("[TVM-FFI ON ] Create args", create_arguments)
compilation_on, compiled = benchmark(
    "[TVM-FFI ON ] Compile operator",
    lambda i: operator.compile(args[i]),
    warmup_it=2,
    profiling_it=5,
)
compiled_artifact = compiled[0]
run_on, _ = benchmark(
    "[TVM-FFI ON ] Launch host overhead ",
    lambda i: operator.run(
        args[i],
        compiled_artifact=compiled_artifact,
        assume_supported_args=True,
        stream=stream,
    ),
)

python

ops.GlobalOptions().use_tvm_ffi = False
args_creation_off, args = benchmark("[TVM-FFI OFF ] Create args", create_arguments)
compilation_off, compiled = benchmark(
    "[TVM-FFI OFF ] Compile operator",
    lambda i: operator.compile(args[i]),
    warmup_it=2,
    profiling_it=5,
)
compiled_artifact = compiled[0]
run_off, _ = benchmark(
    "[TVM-FFI OFF ] Launch host overhead ",
    lambda i: operator.run(
        args[i],
        compiled_artifact=compiled_artifact,
        assume_supported_args=True,
        stream=stream,
    ),
)

# Restore original setting
ops.GlobalOptions().use_tvm_ffi = original_use_tvm_ffi

python

print("Speedups with TVM-FFI: ")
print(f"Arg creation: {args_creation_off / args_creation_on:.2f}x")
print(f"Compilation: {compilation_off / compilation_on:.2f}x")
print(f"Launch host overhead: {run_off / run_on:.2f}x")