Note

This page is a rendered version of 007_heuristics.ipynb available on GitHub.

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.

[1]:

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.

[2]:

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)

Returned 5 operator(s) (limit=5)

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).

[3]:

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

Unsupported GPU: nvmatmul only supports SM100 today; 'H100_SXM' isn't a recognized/supported device (e.g. "B200", "GB200_NVL", "GB300_NVL").

Comparing the result#

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

[4]:

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})")

892 total operator(s) for this problem

After applying heuristics, this was sorted and pruned to 26 operator(s).

fastest heuristic-recommended operator: 0.0894 ms        (name: cutedsl.PersistentDenseGemmOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_2cta_cluster2x2x1_tile256x256x64_schedulerCLC_tma_store)
slowest heuristic-recommended operator: 0.2378 ms        (name: cutedsl.PersistentDenseGemmEFCOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_1cta_cluster1x2x1_tile64x32x64_tma_store)
arbitrary operator: 0.2745 ms    (name: cutedsl.PersistentDenseGemmOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_2cta_cluster2x1x1_tile128x32x64_schedulerCLC_tma_store)

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.