ContextQMD
Back to Cutlass

Heuristics: picking a likely-fast operator

operators/examples/007_heuristics.ipynb

4.8.0dev5.0 KB
Original Source

Heuristics: picking a likely-fast operator

get_operators can return many operators that are all correct for your problem. Without a heuristic, that list is in discovery order — not "fastest first."

A heuristic can optionally rank the estimated performance of candidate operators, and reorders and prunes the candidates accordingly.

CUTLASS Operator API natively supports NVIDIA Matmul Heuristics (nvMatmulHeuristics) to do such ranking directly in get_operators()

Prerequisite: Install nvidia-matmul-heuristics

NVIDIA Matmul Heuristics (nvMatmulHeuristics) is an optional dependency of CUTLASS Operator API, and must be present to run heuristics.

Install: pip install 'nvidia-cutlass-operators[heuristics]'

Currently, heuristics integration is only supported for non-blockscaled GEMMs on SM100.

If you request a heuristic and ranking cannot run (missing optional package, unsupported args), an error is raised.

python
import sys

import torch

import cutlass.operators as ops
from cutlass.operators.heuristics.nvmatmul import is_available as nvmatmul_available

if not (status := ops.utils.device.device_or_env_supports("100")):
    print(f"This notebook expects an SM100-class GPU.\n{status.error}")
    sys.exit(0)

if not nvmatmul_available():
    print(
        "nvmatmul heuristic is unavailable. Install with:\n"
        "  pip install 'nvidia-cutlass-operators[heuristics]'"
    )
    sys.exit(0)

Minimal usage

Operator API can create a nvMatmulHeuristics instance using ops.get_heuristic("nvmatmul")(gpu="B200"). gpu= names the exact GPU SKU to model and defaults to "B200" if omitted.

get_operators() then natively supports using this heuristic returning candidate operators for given arguments. Additionally, we can use limit=N to limit the results to the top-N heuristic-recommended Operators.

Currently, it supports only supports non-blockscaled GEMMs for Blackwell (SM100) GPUs.

When supported, get_operators(args, heuristic) will:

  1. Query the supported, unsorted candidate operators
  2. Query the heuristic for recommended kernel configurations for your given argument (problem size, dtypes, ...).
  3. Matches those recommendations to the candidate Operators. Operators not matched by the heuristic are excluded.
python
M, N, K = 4096, 4096, 4096
A = torch.randn(M, K, device="cuda", dtype=torch.float16)
B = torch.randn(K, N, device="cuda", dtype=torch.float16)
out = torch.empty(M, N, device="cuda", dtype=torch.float16)
args = ops.GemmArguments(A, B, out, accumulator_type=torch.float32)

heuristic = ops.get_heuristic("nvmatmul")(gpu="B200")
# or, equivalently:
heuristic =  ops.heuristics.NvMatmulHeuristics(gpu="B200")

operators = ops.get_operators(
    args,
    target_sm="100a",
    providers=[ops.CuTeDSLProvider],
    heuristic=heuristic,
    limit=5,
)
print(f"Returned {len(operators)} operator(s) (limit=5)")
operators[0].run(args)

An error is raised if the ranking cannot run at all, e.g. for missing package, unsupported args, and unsupported GPU (currently only SM100 is supported).

python
try:
    ops.get_heuristic("nvmatmul")(gpu="H100_SXM")
except ValueError as e:
    print(f"Unsupported GPU: {e}")

Comparing the result

We do a quick benchmark below to compare the operators returned by heuristics.

python
def benchmark_operator(op: ops.Operator, args: ops.GemmArguments, warmup=10, iters=50):
    """Return the median GPU time (in ms) for `op.run(args)`."""
    compiled = op.compile(args)
    for _ in range(warmup):
        op.run(args, compiled_artifact=compiled, assume_supported_args=True)
    torch.cuda.synchronize()

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    times_ms = []
    for _ in range(iters):
        start.record()
        op.run(args, compiled_artifact=compiled, assume_supported_args=True)
        end.record()
        torch.cuda.synchronize()
        times_ms.append(start.elapsed_time(end))
    return sorted(times_ms)[len(times_ms) // 2]


unsorted_operators = ops.get_operators(args, target_sm="100a", providers=[ops.providers.CuTeDSLProvider])
sorted_operators = ops.get_operators(args, target_sm="100a", providers=[ops.providers.CuTeDSLProvider], heuristic=heuristic)


print(f"{len(unsorted_operators)} total operator(s) for this problem\n")
print(f"After applying heuristics, this was sorted and pruned to {len(sorted_operators)} operator(s).\n")

candidates = {
    "fastest heuristic-recommended operator": sorted_operators[0],
    "slowest heuristic-recommended operator": sorted_operators[-1],
    "arbitrary operator": unsorted_operators[0],
}


for label, op in candidates.items():
    median_ms = benchmark_operator(op, args)
    print(f"{label}: {median_ms:.4f} ms \t (name: {op.metadata.operator_name})")

Please note that heuristics-based ranking is an estimate. It is helpful to limit and guide the search to promising candidates to top-N recommendations, and may not yield a definitive single winner.