[Perf] Choose the RDNA3 GEMM tile from the shape

rdna3_f16_gemm builds whatever tile it is handed and defaults to
128x128x32. That tile is right once the problem fills the grid, but it cuts
only 4 workgroups at 256x256 and 16 at 512x512, so on a 96-CU part most CUs
idle no matter how good the inner loop is. Choosing the tile from the shape
is worth up to 3.0x there.

rdna3_f16_gemm_autotune owns that decision in two layers. pick_tile is a
heuristic fitted to a sweep of every feasible tile on 27 shapes; it needs no
GPU and no measurement, and it is what a call resolves to with nothing
configured, so the wrapper benchmarks nothing by default. Above it sits the
shared autotuner: FLYDSL_AUTOTUNE=1 sweeps feasible_tiles for real, and the
result can be frozen into an offline artifact.

The heuristic defaults to 64x64x64 rather than the widest tile that covers
the machine. Measured on gfx1100 it is fastest on 16 of the 27 shapes and
holds 50-59 TFLOP/s throughout, where 128x128x32 swings between 40 and 72.
Taking the widest covering tile cost up to 37% and averaged 6.5%; against
the per-shape fastest tile this averages 0.6%, worst case 8.1%.

Two limits worth knowing. NUM_CU is hard-coded for gfx1100, so the
thresholds do not transfer to a gfx11 part with a different CU count, and
shapes outside the fitted set are extrapolation -- the search exists for
both cases. And _graph_bench, which captures a CUDA graph to get under the
~90us launch overhead that would otherwise swamp these kernels, still reads
the multi-wave tiles a few us high below about 50us, so a tuned result for a
short kernel is a hypothesis to confirm rather than a fact.

feasible_tiles doubles as the search space: anything it excludes does not
divide the shape, cannot fill the prefetch pipeline, or does not fit in LDS,
so benchmarking it would only measure a build failure.

Points the gfx11 benchmark path at the wrapper so its numbers reflect the
chosen tile rather than the default.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: vlluvia

kernels/gemm/rdna3_f16_gemm_autotune.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+"""Tile selection for the RDNA3 WMMA GEMM.
+
+``rdna3_f16_gemm`` builds whatever tile it is handed and defaults to 128x128x32.
+That tile is right once the problem fills the grid, but it cuts only 4
+workgroups at 256x256 and 16 at 512x512, so on a 96-CU part most CUs idle no
+matter how good the inner loop is. Choosing the tile from the shape is worth up
+to 3.0x there, and this module owns that decision in two layers:
+
+  * ``pick_tile`` — a heuristic fitted to a sweep of every feasible tile on 27
+    shapes. It needs no GPU and no measurement, and it is what a call resolves
+    to with nothing configured, so the wrapper benchmarks nothing by default.
+  * the shared autotuner — ``FLYDSL_AUTOTUNE=1`` sweeps ``feasible_tiles`` for
+    real on the GPU in hand, and the result can be frozen into an offline
+    artifact so other machines with the same device fingerprint skip the search.
+
+The second layer earns its keep mainly where the first cannot reach: ``NUM_CU``
+is hard-coded for gfx1100, so the thresholds do not transfer to a gfx11 part
+with a different CU count, and shapes outside the fitted set are extrapolation.
+It is also the least settled part of this — see ``_graph_bench`` for why its
+verdict on the shortest kernels should not be taken at face value. The default
+path does not benchmark and is unaffected.
+
+The tile is not a Constexpr argument of one compiled entry point: it decides the
+block shape, the wave grid and the LDS budget, so each candidate is a separate
+module. The dispatcher below therefore takes the tile as ordinary keyword
+arguments and looks the built module up in a cache, the same shape of
+indirection ``conv3d_implicit_autotune`` uses.
+"""
+
+import functools
+
+import torch
+
+from flydsl.autotune import Config, autotune, do_bench
+from kernels.gemm.rdna3_f16_gemm import K_PAD, WAVE_SIZE, WMMA_K, WMMA_M, WMMA_N, create_wmma_gemm_module
+
+# gfx1100 (W7900) exposes 96 CUs. Used only to decide when a tile is too coarse
+# to fill the machine; being off by a little just shifts one ladder step.
+NUM_CU = 96
+
+# LDS budget per workgroup. K_PAD comes from the kernel so the feasibility check
+# below cannot drift from the allocation it is predicting.
+LDS_BYTES = 64 * 1024
+
+# Named by the block tile they produce: (reg_m, reg_n, reg_k, waves_m, waves_n).
+TILE_128x128x32 = (4, 4, 2, 2, 2)
+TILE_128x64x32 = (4, 2, 2, 2, 2)
+TILE_64x64x64 = (2, 2, 4, 2, 2)
+TILE_32x64x64 = (2, 2, 4, 1, 2)
+TILE_32x32x64 = (2, 2, 4, 1, 1)
+
+# Tile ladder, widest first.
+#
+# 128x64x32 exists only on the small-K ladder: with large K a workgroup runs long
+# enough that the deeper k-tile (fewer barriers, twice as long for the gmem
+# prefetch to land) pays off, while with small K the per-workgroup prologue and
+# epilogue dominate and the wider tile amortizes them.
+#
+# The ladder is the feasibility-ordered search space; pick_tile does not walk it
+# in order. 32x64x64 is here only as a fallback for shapes the others cannot
+# divide -- it was not the fastest tile on any of the 27 shapes swept.
+_LADDER_LARGE_K = [TILE_128x128x32, TILE_64x64x64, TILE_32x64x64, TILE_32x32x64]
+_LADDER_SMALL_K = [TILE_128x128x32, TILE_128x64x32, TILE_64x64x64, TILE_32x64x64, TILE_32x32x64]
+
+
+def _tile_workgroups(M, N, K, cfg):
+    """Workgroup count for this tile, or None if the shape cannot use it."""
+    reg_m, reg_n, reg_k, waves_m, waves_n = cfg
+    block_m = WMMA_M * reg_m * waves_m
+    block_n = WMMA_N * reg_n * waves_n
+    block_k = WMMA_K * reg_k
+    threads = waves_m * waves_n * WAVE_SIZE
+    if M % block_m or N % block_n or K % block_k:
+        return None
+    if K // block_k < 2:  # the prefetch pipeline needs at least two k-tiles
+        return None
+    # Every thread must carry a whole 8-element vector of both tiles.
+    if (block_m * block_k) % (threads * 8) or (block_n * block_k) % (threads * 8):
+        return None
+    if 2 * (block_m + block_n) * (block_k + K_PAD) * 2 > LDS_BYTES:  # 2 buffers, 2 bytes/elem
+        return None
+    return (M // block_m) * (N // block_n)
+
+
+def _ladder_for(K):
+    return _LADDER_SMALL_K if K <= 1024 else _LADDER_LARGE_K
+
+
+def feasible_tiles(M, N, K):
+    """``(tile, workgroup count)`` for every ladder tile this shape can run, widest first.
+
+    Also the search space the autotuner sweeps: anything excluded here does not
+    divide the shape, cannot fill the prefetch pipeline, or does not fit in LDS,
+    so benchmarking it would only measure a build failure.
+    """
+    return [(cfg, wgs) for cfg in _ladder_for(K) if (wgs := _tile_workgroups(M, N, K, cfg)) is not None]
+
+
+def pick_tile(M, N, K):
+    """Tile for this shape, fitted to a sweep of every ladder tile on 27 shapes.
+
+    64x64x64 is the default rather than the widest tile that covers the machine.
+    Measured on gfx1100 it is the fastest tile on 16 of the 27 shapes and holds
+    50-59 TFLOP/s across the whole range, where 128x128x32 swings between 40 and
+    72 depending on how its much coarser grid happens to land. Taking the widest
+    covering tile cost up to 37% (1664x1664x1024) and averaged 6.5%.
+
+    Three exceptions, in order:
+
+      * 128x128x32 once the grid is worth at least ~2.5 workgroups per CU. Its
+        compute intensity wins outright there, by 5-11% over 64x64x64.
+      * 128x64x32 (small-K ladder only) when it lands near one workgroup per CU.
+        That is the narrow band around 1024x1024, where it leads by 13-28%.
+      * 32x32x64 when 64x64x64 cannot fill a quarter of the machine and K is
+        long enough for the idle CUs to dominate: 256x256x4096, worth 23%.
+
+    Against the per-shape fastest tile this averages 0.6%, worst case 8.1% at
+    1152x1152x1024, where 128x128x32's grid happens to land well.
+    """
+    feasible = dict(feasible_tiles(M, N, K))
+    if not feasible:
+        return _ladder_for(K)[0]
+
+    if feasible.get(TILE_128x128x32, 0) >= 2.5 * NUM_CU:
+        return TILE_128x128x32
+    if NUM_CU <= feasible.get(TILE_128x64x32, 0) <= 1.5 * NUM_CU:
+        return TILE_128x64x32
+    if TILE_64x64x64 in feasible:
+        starved = feasible[TILE_64x64x64] < NUM_CU / 4 and K >= 2048
+        if starved and TILE_32x32x64 in feasible:
+            return TILE_32x32x64
+        return TILE_64x64x64
+    return list(feasible)[-1]
+
+
+# Launches to capture per graph. Enough that replay is dominated by the kernel
+# rather than by graph launch, small enough to keep capture cheap.
+_GRAPH_LAUNCHES = 20
+_REPLAYS_PER_ROUND = 8
+
+_TILE_FIELDS = ("reg_m", "reg_n", "reg_k", "waves_m", "waves_n")
+
+# Launcher for a call signature whose tile the autotuner has already resolved.
+# The tuner re-derives its cache key from scratch on every call — fingerprinting
+# the environment, toolchain and device — which costs more host time than a
+# small GEMM takes on the GPU, and under FLYDSL_AUTOTUNE=1 it re-runs the whole
+# search per call. Consulting it once per signature keeps steady-state dispatch
+# as cheap as calling the built module directly.
+_resolved = {}
+
+
+def _tile_config(tile):
+    return Config(**dict(zip(_TILE_FIELDS, tile)))
+
+
+@functools.lru_cache(maxsize=None)
+def _build(M, N, K, in_dtype, out_dtype, rounding, reg_m, reg_n, reg_k, waves_m, waves_n):
+    launch_fn, _, _, _ = create_wmma_gemm_module(
+        M,
+        N,
+        K,
+        in_dtype=in_dtype,
+        out_dtype=out_dtype,
+        rounding=rounding,
+        reg_m=reg_m,
+        reg_n=reg_n,
+        reg_k=reg_k,
+        waves_m=waves_m,
+        waves_n=waves_n,
+    )
+    return launch_fn
+
+
+def rdna3_gemm_dispatch(
+    C,
+    A,
+    B_T,
+    M,
+    N,
+    K,
+    in_dtype="bf16",
+    out_dtype="bf16",
+    rounding="rn",
+    reg_m=None,
+    reg_n=None,
+    reg_k=None,
+    waves_m=None,
+    waves_n=None,
+    stream=None,
+    sr_seed=0,
+):
+    """Run the GEMM on one tile. Unset tile fields fall back to ``pick_tile``.
+
+    The stream is resolved here rather than by the caller so that this stays
+    capturable: under ``torch.cuda.graph`` the current stream is the capture
+    stream, and enqueueing onto a stream captured before then aborts the capture.
+    """
+    if stream is None:
+        stream = torch.cuda.current_stream()
+    # Resolve before the cache key so a partially specified tile and the fully
+    # spelled-out one it means share a single built module.
+    tile = tuple(
+        auto if given is None else given
+        for auto, given in zip(pick_tile(M, N, K), (reg_m, reg_n, reg_k, waves_m, waves_n))
+    )
+    launch_fn = _build(M, N, K, in_dtype, out_dtype, rounding, *tile)
+    # A search calls this once per candidate and then once more on the winner,
+    # so the last write is the config the tuner settled on.
+    _resolved[(M, N, K, in_dtype, out_dtype, rounding)] = launch_fn
+    return launch_fn(C, A, B_T, stream, sr_seed)
+
+
+def _default_config(
+    C=None,
+    A=None,
+    B_T=None,
+    M=None,
+    N=None,
+    K=None,
+    in_dtype="bf16",
+    out_dtype="bf16",
+    rounding="rn",
+    **_kwargs,
+):
+    return _tile_config(pick_tile(M, N, K))
+
+
+def _search_configs(
+    C=None,
+    A=None,
+    B_T=None,
+    M=None,
+    N=None,
+    K=None,
+    in_dtype="bf16",
+    out_dtype="bf16",
+    rounding="rn",
+    **_kwargs,
+):
+    candidates = [_tile_config(tile) for tile, _wgs in feasible_tiles(M, N, K)]
+    return candidates or [_default_config(M=M, N=N, K=K)]
+
+
+def _graph_bench(fn, warmup=5, rep=25):
+    """Fastest observed ms per launch, timed by replaying a captured graph.
+
+    The stock ``do_bench`` pays one launch plus one full sync per measurement,
+    about 90us of host time on this kernel. That is longer than the kernel runs
+    on any shape small enough for the tile to be worth choosing, so all the
+    candidates measure alike and the search ends up ranking dispatch noise.
+    Capturing the launches amortises that overhead away and makes a sweep
+    reproducible to within a percent.
+
+    It is still not trustworthy on the shortest kernels. Measured in isolation,
+    512x512x2048 runs at 28.6us on 64x64x64 and 31.0us on 32x32x64; measured
+    here the multi-wave tiles read about 5us high and the ranking inverts, so a
+    forced sweep of that shape emits an artifact for the slower tile. Treat a
+    tuned result for a sub-50us kernel as a hypothesis to confirm, not a fact.
+
+    Falls back to the stock timer if the kernel turns out not to be capturable.
+    """
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    try:
+        with torch.cuda.graph(graph):
+            for _ in range(_GRAPH_LAUNCHES):
+                fn()
+    except Exception:
+        return do_bench(fn, warmup=warmup, rep=rep)
+
+    for _ in range(3):
+        graph.replay()
+    torch.cuda.synchronize()
+
+    # Time a batch of replays per round so the graph launch is amortised too,
+    # and keep the fastest round. This runs on shared nodes where a neighbour
+    # can inflate a whole round two- or threefold; taking the median carries
+    # those rounds into the result, taking the minimum drops them.
+    rounds = max(3, rep // _REPLAYS_PER_ROUND)
+    best = float("inf")
+    for _ in range(rounds):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        for _ in range(_REPLAYS_PER_ROUND):
+            graph.replay()
+        end.record()
+        torch.cuda.synchronize()
+        best = min(best, start.elapsed_time(end) / (_REPLAYS_PER_ROUND * _GRAPH_LAUNCHES))
+    return best
+
+
+_gemm_tuner = autotune(
+    configs=_search_configs,
+    key=["M", "N", "K", "in_dtype", "out_dtype", "rounding"],
+    default=_default_config,
+    do_bench=_graph_bench,
+    artifact_name="rdna3_f16_gemm",
+)(rdna3_gemm_dispatch)
+
+
+def rdna3_gemm_autotuned(
+    C,
+    A,
+    B_T,
+    in_dtype="bf16",
+    out_dtype="bf16",
+    rounding="rn",
+    stream=None,
+    sr_seed=0,
+):
+    """``C = A @ B_T.T`` on the tile chosen for this shape."""
+    M, K = A.shape
+    N = B_T.shape[0]
+    M, N, K = int(M), int(N), int(K)
+
+    launch_fn = _resolved.get((M, N, K, in_dtype, out_dtype, rounding))
+    if launch_fn is not None:
+        return launch_fn(C, A, B_T, torch.cuda.current_stream() if stream is None else stream, sr_seed)
+
+    # Make the caller's stream current instead of passing it down, so the
+    # dispatcher picks it up while a benchmark capture still gets its own.
+    launch_stream = torch.cuda.current_stream() if stream is None else stream
+    with torch.cuda.device(A.device), torch.cuda.stream(launch_stream):
+        return _gemm_tuner(
+            C,
+            A,
+            B_T,
+            M=M,
+            N=N,
+            K=K,
+            in_dtype=in_dtype,
+            out_dtype=out_dtype,
+            rounding=rounding,
+            stream=None,
+            sr_seed=sr_seed,
+        )

tests/kernels/benchmark_common.py

@@ -239,22 +239,31 @@ def _bench_flydsl_torch(*, op: str, M: int, N: int, dtype: str, warmup: int, ite
         # arch.
         from flydsl.runtime.device import get_rocm_arch as _get_arch
 
-        if str(_get_arch() or "").startswith("gfx11"):
-            from kernels.gemm.rdna3_f16_gemm import create_wmma_gemm_module
-        else:
-            from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module
-
         K = N  # square by default; caller can override via config
         torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16
-        launch, *_ = create_wmma_gemm_module(M, N, K, in_dtype=dtype, out_dtype="bf16")
         A = torch.randn(M, K, dtype=torch_dtype, device="cuda")
         B_T = torch.randn(N, K, dtype=torch_dtype, device="cuda")
         C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda")
-        return bench_gpu_us_torch(
-            lambda: launch(C, A, B_T, torch.cuda.current_stream()),
-            warmup=warmup,
-            iters=iters,
-        )
+
+        if str(_get_arch() or "").startswith("gfx11"):
+            # Through the autotune wrapper, so the number reflects the tile
+            # chosen for the shape rather than the kernel's 128x128x32 default.
+            # With nothing configured that resolution still benchmarks nothing,
+            # and warmup absorbs the one-off dispatch through the tuner.
+            from kernels.gemm.rdna3_f16_gemm_autotune import rdna3_gemm_autotuned
+
+            def run():
+                rdna3_gemm_autotuned(C, A, B_T, in_dtype=dtype, out_dtype="bf16")
+
+        else:
+            from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module
+
+            launch, *_ = create_wmma_gemm_module(M, N, K, in_dtype=dtype, out_dtype="bf16")
+
+            def run():
+                launch(C, A, B_T, torch.cuda.current_stream())
+
+        return bench_gpu_us_torch(run, warmup=warmup, iters=iters)
 
     if op == "wmma_fp8_gemm":
         from kernels.gemm.rdna_fp8_preshuffle_gemm import (