From 6c8e7c31e9701f86fc767e86a1c2d2c00feb3cfc Mon Sep 17 00:00:00 2001 From: vlluvia Date: Fri, 7 Aug 2026 11:13:17 +0800 Subject: [PATCH 1/2] [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 --- kernels/gemm/rdna3_f16_gemm_autotune.py | 343 ++++++++++++++++++++++++ tests/kernels/benchmark_common.py | 31 ++- tests/kernels/test_rdna_gemm.py | 42 +++ tests/unit/test_rdna3_gemm_autotune.py | 89 ++++++ tests/unit/test_rdna3_tile_selection.py | 161 +++++++++++ 5 files changed, 655 insertions(+), 11 deletions(-) create mode 100644 kernels/gemm/rdna3_f16_gemm_autotune.py create mode 100644 tests/unit/test_rdna3_gemm_autotune.py create mode 100644 tests/unit/test_rdna3_tile_selection.py diff --git a/kernels/gemm/rdna3_f16_gemm_autotune.py b/kernels/gemm/rdna3_f16_gemm_autotune.py new file mode 100644 index 000000000..fe2336de4 --- /dev/null +++ b/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, + ) diff --git a/tests/kernels/benchmark_common.py b/tests/kernels/benchmark_common.py index 8770f3915..2c02230f2 100644 --- a/tests/kernels/benchmark_common.py +++ b/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 ( diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 77aab2beb..c48c9ba40 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/tests/kernels/test_rdna_gemm.py @@ -21,6 +21,7 @@ from flydsl.runtime.device import get_rocm_arch # noqa: E402 from kernels.gemm.rdna3_f16_gemm import create_wmma_gemm_module as _create_wmma_gemm_module_gfx11 # noqa: E402 +from kernels.gemm.rdna3_f16_gemm_autotune import pick_tile # noqa: E402 from kernels.gemm.rdna_f16_gemm import create_wmma_gemm_module as _create_wmma_gemm_module_gfx12 # noqa: E402 from kernels.gemm.rdna_fp8_preshuffle_gemm import ( # noqa: E402 compile_fp8_gemm, @@ -209,6 +210,47 @@ def test_f16_gemm_grid_m_not_a_multiple_of_the_group_width(M, N, K): assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05) +@pytest.mark.parametrize( + "M, N, K", + [ + pytest.param(256, 256, 4096, id="256x256x4096"), + pytest.param(1024, 1024, 1024, id="1024x1024x1024"), + ], +) +def test_f16_gemm_autotuned_matches_the_heuristic_path(M, N, K): + """The autotune wrapper, left unconfigured, is a pass-through. + + It resolves through the tuner once and then calls the built module directly: + the tuner re-derives its cache key on every call, which costs more host time + than these shapes take on the GPU. So this checks both halves — the same + tile as ``pick_tile``, and that the second call does not build again. + """ + _requires_rdna3() + from kernels.gemm import rdna3_f16_gemm_autotune as gemm_autotune + + torch.manual_seed(42) + A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1 + B_T = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1 + C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda") + + signature = (M, N, K, "bf16", "bf16", "rn") + gemm_autotune._resolved.pop(signature, None) + + gemm_autotune.rdna3_gemm_autotuned(C, A, B_T) + torch.cuda.synchronize() + assert verify_output(C.float(), A.float() @ B_T.float().T, atol=0.05, rtol=0.05) + + resolved = gemm_autotune._resolved[signature] + expected_tile = pick_tile(M, N, K) + assert resolved is gemm_autotune._build(M, N, K, "bf16", "bf16", "rn", *expected_tile) + + C.zero_() + gemm_autotune.rdna3_gemm_autotuned(C, A, B_T) + torch.cuda.synchronize() + assert verify_output(C.float(), A.float() @ B_T.float().T, atol=0.05, rtol=0.05) + assert gemm_autotune._resolved[signature] is resolved + + @pytest.mark.parametrize( "M, N, K", [ diff --git a/tests/unit/test_rdna3_gemm_autotune.py b/tests/unit/test_rdna3_gemm_autotune.py new file mode 100644 index 000000000..3988f0cac --- /dev/null +++ b/tests/unit/test_rdna3_gemm_autotune.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors +"""GPU-free tests for the RDNA3 GEMM autotune wrapper. + +The wrapper's job is to expose tile selection through the shared autotuner +without a call having to opt into benchmarking, so the properties worth pinning +are about agreement rather than performance: + + * with nothing configured, the tuner's ``default`` is exactly ``pick_tile`` — + a call that never benchmarks must land on the heuristic's tile; + * the search space only contains tiles the shape can actually build, and it + always contains the default, so a search can confirm the heuristic and can + never be forced to pick something that fails at build time; + * ``feasible_tiles`` reports the same tiles, in the same order, that + ``pick_tile`` chooses from. + +Device-side coverage is ``tests/kernels/test_rdna_gemm.py``. +""" + +import pytest + +from kernels.gemm.rdna3_f16_gemm_autotune import ( + _TILE_FIELDS, + _default_config, + _ladder_for, + _search_configs, + _tile_workgroups, + feasible_tiles, + pick_tile, +) + +pytestmark = pytest.mark.l0_backend_agnostic + +SHAPES = [ + (M, N, K) + for M in (256, 512, 768, 1024, 1536, 2048, 4096) + for N in (256, 512, 1024, 2048) + for K in (512, 1024, 2048, 4096) +] + +_shape_id = lambda shape: "x".join(map(str, shape)) # noqa: E731 + + +def _tile_of(config): + return tuple(config.kwargs[field] for field in _TILE_FIELDS) + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_default_config_is_the_heuristic_tile(shape): + """An untuned call must be indistinguishable from calling the kernel directly.""" + M, N, K = shape + assert _tile_of(_default_config(M=M, N=N, K=K)) == pick_tile(*shape) + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_search_space_only_holds_buildable_tiles(shape): + """Every candidate must survive ``_tile_workgroups``. + + A candidate the shape cannot use does not merely lose the benchmark: it + raises out of ``create_wmma_gemm_module``, so the search would spend its + time measuring build failures. + """ + M, N, K = shape + if not feasible_tiles(*shape): + pytest.skip("no tile in the ladder fits this shape") + for config in _search_configs(M=M, N=N, K=K): + assert _tile_workgroups(*shape, _tile_of(config)) is not None + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_default_is_reachable_by_the_search(shape): + """The search must be able to return the heuristic's own answer.""" + M, N, K = shape + if not feasible_tiles(*shape): + pytest.skip("no tile in the ladder fits this shape") + tiles = [_tile_of(config) for config in _search_configs(M=M, N=N, K=K)] + assert pick_tile(*shape) in tiles + + +@pytest.mark.parametrize("shape", SHAPES, ids=_shape_id) +def test_feasible_tiles_is_the_ladder_filtered_in_order(shape): + """``pick_tile`` walks this list, so it has to stay a subsequence of the ladder.""" + ladder = _ladder_for(shape[2]) + reported = feasible_tiles(*shape) + + assert [tile for tile, _ in reported] == [cfg for cfg in ladder if _tile_workgroups(*shape, cfg) is not None] + for tile, workgroups in reported: + assert workgroups == _tile_workgroups(*shape, tile) diff --git a/tests/unit/test_rdna3_tile_selection.py b/tests/unit/test_rdna3_tile_selection.py new file mode 100644 index 000000000..cfca1acac --- /dev/null +++ b/tests/unit/test_rdna3_tile_selection.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors +"""GPU-free tests for the RDNA3 GEMM host-side tile selection. + +``pick_tile`` decides which tile a shape gets. Beyond the measured expectations, +two properties make it safe to enable by default: it never returns a tile the +shape cannot build, and it never moves a shape off 128x128x32 while that tile +still has a deep enough grid to win. + +All plain integer arithmetic that runs before a kernel is built, so no GPU is +needed. The swizzle the chosen tile feeds into is covered by +``test_rdna3_grid_swizzle.py``, and the device-side counterpart is +``tests/kernels/test_rdna_gemm.py``. +""" + +import pytest + +from kernels.gemm.rdna3_f16_gemm import _group_width, _swizzle_tile_id +from kernels.gemm.rdna3_f16_gemm_autotune import ( + NUM_CU, + TILE_32x32x64, + TILE_32x64x64, + TILE_64x64x64, + TILE_128x64x32, + TILE_128x128x32, + _ladder_for, + _tile_workgroups, + feasible_tiles, + pick_tile, +) + +pytestmark = pytest.mark.l0_backend_agnostic + +DEFAULT_TILE = TILE_128x128x32 # the tile used before selection existed +DEFAULT_GROUP_M = 8 # create_wmma_gemm_module's default grouping cap + +# Every shape the tile heuristic was measured on, with the tile it is expected to +# pick, so that a heuristic edit has to state which shapes it moves. Each of +# these was timed against the whole ladder; the expectation is the fastest tile +# except where noted. +MEASURED_SHAPES = [ + pytest.param((256, 256, 4096), TILE_32x32x64, id="256x256x4096-32x32x64"), + pytest.param((256, 256, 1024), TILE_64x64x64, id="256x256x1024-64x64x64"), + pytest.param((384, 384, 2048), TILE_64x64x64, id="384x384x2048-64x64x64"), + pytest.param((512, 512, 1024), TILE_64x64x64, id="512x512x1024-64x64x64"), + pytest.param((512, 512, 4096), TILE_64x64x64, id="512x512x4096-64x64x64"), + pytest.param((256, 1024, 4096), TILE_64x64x64, id="256x1024x4096-64x64x64"), + pytest.param((768, 768, 2048), TILE_64x64x64, id="768x768x2048-64x64x64"), + pytest.param((1024, 1024, 512), TILE_128x64x32, id="1024x1024x512-128x64x32"), + pytest.param((1024, 1024, 1024), TILE_128x64x32, id="1024x1024x1024-128x64x32"), + pytest.param((1024, 1024, 4096), TILE_64x64x64, id="1024x1024x4096-64x64x64"), + # 128x128x32 is 8.1% faster here, the worst case the heuristic accepts. + pytest.param((1152, 1152, 1024), TILE_64x64x64, id="1152x1152x1024-64x64x64"), + pytest.param((1536, 1536, 1024), TILE_64x64x64, id="1536x1536x1024-64x64x64"), + pytest.param((1792, 1792, 1024), TILE_64x64x64, id="1792x1792x1024-64x64x64"), + pytest.param((2048, 2048, 512), DEFAULT_TILE, id="2048x2048x512-128x128x32"), + pytest.param((2048, 2048, 2048), DEFAULT_TILE, id="2048x2048x2048-128x128x32"), + pytest.param((3072, 3072, 1024), DEFAULT_TILE, id="3072x3072x1024-128x128x32"), + pytest.param((4096, 4096, 4096), DEFAULT_TILE, id="4096x4096x4096-128x128x32"), +] + +MEASURED_ONLY_SHAPES = [p.values[0] for p in MEASURED_SHAPES] + +# A spread for the property tests: square and skewed, K on both sides of the +# ladder split, sizes from "cannot fill the machine" up to "fills it easily". +PROPERTY_SHAPES = [ + (M, N, K) + for M in (256, 512, 768, 1024, 1536, 2048, 4096) + for N in (256, 512, 1024, 2048) + for K in (512, 1024, 2048, 4096) +] + +_shape_id = lambda shape: "x".join(map(str, shape)) # noqa: E731 + + +def _block_shape(cfg): + reg_m, reg_n, reg_k, waves_m, waves_n = cfg + return 16 * reg_m * waves_m, 16 * reg_n * waves_n, 16 * reg_k + + +def _tile_grid(shape): + """(grid_m, grid_n) the kernel would launch for this shape.""" + M, N, _ = shape + block_m, block_n, _ = _block_shape(pick_tile(*shape)) + return M // block_m, N // block_n + + +@pytest.mark.parametrize("shape, expected", MEASURED_SHAPES) +def test_pick_tile_matches_measured_shapes(shape, expected): + """The tile picked for each benchmarked shape is the one that was measured.""" + assert pick_tile(*shape) == expected + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_picked_tile_is_usable_by_the_shape(shape): + """Whenever the ladder has a usable tile, the picked one is usable. + + ``create_wmma_gemm_module`` asserts the divisibility and vectorization rules + that ``_tile_workgroups`` screens for, so returning an unusable tile turns + into an assertion failure at build time. + """ + if all(_tile_workgroups(*shape, cfg) is None for cfg in _ladder_for(shape[2])): + pytest.skip("no tile in the ladder fits this shape") + assert _tile_workgroups(*shape, pick_tile(*shape)) is not None + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_deep_grids_keep_the_original_tile(shape): + """A shape whose 128x128x32 grid is deep enough must not be moved off it. + + This is what makes the heuristic safe to turn on by default: the large + shapes take the same path as before selection existed, so they cannot + regress. Merely covering the machine is not enough — 128x128x32 lost by 19% + at 1792x1792x1024 and 37% at 1664x1664x1024, both well past one workgroup + per CU — so the bar is the measured one. + """ + wgs = _tile_workgroups(*shape, DEFAULT_TILE) + if wgs is None or wgs < 2.5 * NUM_CU: + pytest.skip("128x128x32's grid is not deep enough for this shape") + assert pick_tile(*shape) == DEFAULT_TILE + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_the_narrowest_tile_needs_a_starved_grid(shape): + """32x32x64 trades compute intensity for parallelism, and loses when it does not need to. + + It gave up 12-23% against 64x64x64 wherever 64x64x64 had the workgroups to + keep the machine busy, so it must be unreachable unless 64x64x64 either does + not fit or cannot fill the machine. + """ + if pick_tile(*shape) != TILE_32x32x64: + pytest.skip("this shape does not take the narrowest tile") + feasible = dict(feasible_tiles(*shape)) + assert TILE_64x64x64 not in feasible or feasible[TILE_64x64x64] < NUM_CU / 4 + + +@pytest.mark.parametrize("shape", PROPERTY_SHAPES, ids=_shape_id) +def test_32x64x64_is_only_ever_a_fallback(shape): + """32x64x64 was not the fastest tile on any of the 27 shapes swept. + + The old heuristic chose it for five of them, each time losing to 64x64x64. + It stays on the ladder only to cover shapes that nothing else divides. + """ + if pick_tile(*shape) != TILE_32x64x64: + pytest.skip("this shape does not take 32x64x64") + assert dict(feasible_tiles(*shape)).keys() == {TILE_32x64x64} + + +@pytest.mark.parametrize("shape", MEASURED_ONLY_SHAPES, ids=_shape_id) +def test_picked_tile_and_default_grouping_stay_inside_the_grid(shape): + """Selection and swizzle have to agree, not just work in isolation. + + The fault needs both halves to line up: the tile decides grid_m, and only then + does the grouping width become a divisor question. + """ + grid_m, grid_n = _tile_grid(shape) + width = _group_width(grid_m, DEFAULT_GROUP_M) + mapped = [_swizzle_tile_id(pid, grid_n, width) for pid in range(grid_m * grid_n)] + + assert set(mapped) == {(m, n) for m in range(grid_m) for n in range(grid_n)} From ac1af8cb4fad1715e053f122322eb1bf666d8e46 Mon Sep 17 00:00:00 2001 From: vlluvia Date: Fri, 14 Aug 2026 15:11:00 +0800 Subject: [PATCH 2/2] [Perf] Add wide RDNA3 GEMM tiles Extend the RDNA3 GEMM tile ladder with wide large-shape tiles that fit through an unpadded LDS layout, and pin DeepSeek-V4-Flash decode rows as explicit follow-on scope. --- kernels/gemm/rdna3_f16_gemm.py | 193 ++++++++++++++++++++---- kernels/gemm/rdna3_f16_gemm_autotune.py | 59 +++++++- tests/unit/test_rdna3_gemm_autotune.py | 24 +++ 3 files changed, 243 insertions(+), 33 deletions(-) diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index 024bd9f22..1c7f49de9 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/kernels/gemm/rdna3_f16_gemm.py @@ -50,6 +50,48 @@ K_PAD = 8 +def _sched_plan(reg_m, reg_n, reg_k, g2s_chunks): + """The order the k-tile body is meant to issue in, as (group, count) pairs. + + Left alone, LLVM hoists the whole ds_store block into the middle of the WMMA + stream, so each store's ``s_waitcnt vmcnt`` stalls the wave while the WMMAs + behind it wait. This spends the WMMA stream as cover instead: global loads + go out first, the next k-step's LDS reads hide under this k-step's math, and + the stores drain under the tail, which also leaves the ``lgkmcnt(0)`` in + front of the barrier with almost nothing left to wait for. + + Built here, from Python ints, because the kernel body is traced and cannot + branch on values. The counts are hints: a group the dependences cannot + satisfy is dropped rather than honoured. + """ + per_rk_wmma = reg_m * reg_n + per_rk_dsrd = 2 * (reg_m + reg_n) # each v16 operand is two 128-bit reads + chunk = max(1, reg_n) + + plan = [("vmem", g2s_chunks), ("dsrd", per_rk_dsrd)] + dsrd_left = per_rk_dsrd * (reg_k - 1) + dsrd_step = chunk + dswr_left = g2s_chunks + dswr_chunk = max(1, g2s_chunks // 2) + for _ in range(reg_k): + issued = 0 + while issued < per_rk_wmma: + n = min(chunk, per_rk_wmma - issued) + plan.append(("mfma", n)) + issued += n + if dsrd_left: + take = min(dsrd_step, dsrd_left) + plan.append(("dsrd", take)) + dsrd_left -= take + elif dswr_left: + take = min(dswr_chunk, dswr_left) + plan.append(("dswr", take)) + dswr_left -= take + if dswr_left: + plan.append(("dswr", dswr_left)) + return tuple(plan) + + def _group_width(grid_m, group_m): """Largest grouping width <= group_m that divides grid_m. @@ -95,6 +137,8 @@ def create_wmma_gemm_module( group_m=8, a_k_pad=K_PAD, b_k_pad=K_PAD, + lds_layout="pad", + sched_hint=False, ): gpu_arch = str(get_rocm_arch() or "") if not gpu_arch.startswith("gfx11"): @@ -109,7 +153,10 @@ def create_wmma_gemm_module( NUM_WAVES = waves_m * waves_n # 4 THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE # 128 - assert reg_k >= 2 and reg_k % 2 == 0 + # One WMMA K-step is 16 elements and the v16 operand reads it as two 8-wide + # chunks, so BLOCK_K only has to be a multiple of 16. reg_k=1 (BLOCK_K=16) is + # what keeps the LDS tile small enough to fit more than one workgroup per CU. + assert reg_k >= 1 assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}" if rounding == "rs": assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only" @@ -121,13 +168,37 @@ def create_wmma_gemm_module( # ``row = tid // THRS_K, col = (tid % THRS_K) * LOAD_VEC``. THRS_K = BLOCK_K // LOAD_VEC THRS_M = THREADS_PER_BLOCK // THRS_K + # 128-bit chunks of A and B one thread moves per k-tile; equals both the + # global-load and the ds_store count in the loop body. + G2S_CHUNKS = (BLOCK_M + BLOCK_N) * BLOCK_K // THREADS_PER_BLOCK // LOAD_VEC + SCHED_PLAN = _sched_plan(reg_m, reg_n, reg_k, G2S_CHUNKS) if sched_hint else () assert THRS_K * THRS_M == THREADS_PER_BLOCK assert BLOCK_M % THRS_M == 0 and BLOCK_N % THRS_M == 0 - BLOCK_K_PAD_A = BLOCK_K + a_k_pad # 40 - BLOCK_K_PAD_B = BLOCK_K + b_k_pad # 40 - LDS_A_SIZE = BLOCK_M * BLOCK_K_PAD_A - LDS_B_SIZE = BLOCK_N * BLOCK_K_PAD_B + # Two ways to lay a tile out in LDS, and the choice is what decides which + # macro tiles are reachable at all: + # + # "pad" row-major, BLOCK_K + 8 elements per row. The pad is what keeps + # the 128-bit reads off each other's banks, and it is the only + # pad that does: the row start has to stay 16-byte aligned, which + # leaves 0, 8, 16 and 24, and of those only 8 and 24 spread the + # 16 rows a read touches over all 32 banks. + # + # "kblock" k-major in groups of 8, so element (row, k) sits at + # ((k // 8) * rows + row) * 8 + k % 8. Consecutive rows are then + # 16 bytes apart, so the 8 lanes the LDS services per cycle cover + # 128 contiguous bytes -- every bank, once -- with no pad at all. + # + # Dropping the pad is worth 20% of the LDS budget, which is what brings a + # 256x256x32 tile inside the 64 KB a workgroup may allocate. + assert lds_layout in ("pad", "kblock") + if lds_layout == "kblock": + assert BLOCK_K % LOAD_VEC == 0 + a_k_pad = b_k_pad = 0 + ROW_STRIDE_A = BLOCK_K + a_k_pad + ROW_STRIDE_B = BLOCK_K + b_k_pad + LDS_A_SIZE = BLOCK_M * ROW_STRIDE_A + LDS_B_SIZE = BLOCK_N * ROW_STRIDE_B LDS_ONE_BUF = LDS_A_SIZE + LDS_B_SIZE LDS_TOTAL = 2 * LDS_ONE_BUF @@ -248,14 +319,30 @@ def _v8_load(v8_idx): # allocation instead, which is what the flat _v8_store did by hand. def _lds_dst(buf_offset, base, rows, row_stride): ptr = fx.add_offset(lds_ptr, fx.make_int_tuple(buf_offset + base)) - view = fx.make_view(fx.recast_iter(elem_dtype, ptr), fx.make_layout((rows, BLOCK_K), (row_stride, 1))) + if const_expr(lds_layout == "kblock"): + # (row, k) with k split as (k % 8, k // 8): the low part is + # contiguous so the 128-bit copy atom still sees eight adjacent + # elements, the high part strides a whole plane of rows. + layout = fx.make_layout( + (rows, (LOAD_VEC, BLOCK_K // LOAD_VEC)), + (LOAD_VEC, (1, rows * LOAD_VEC)), + ) + else: + layout = fx.make_layout((rows, BLOCK_K), (row_stride, 1)) + view = fx.make_view(fx.recast_iter(elem_dtype, ptr), layout) return thr_g2s.partition_D(view)[None, None, None] def _pA_s(buf_offset): - return _lds_dst(buf_offset, 0, BLOCK_M, BLOCK_K_PAD_A) + return _lds_dst(buf_offset, 0, BLOCK_M, ROW_STRIDE_A) def _pB_s(buf_offset): - return _lds_dst(buf_offset, LDS_A_SIZE, BLOCK_N, BLOCK_K_PAD_B) + return _lds_dst(buf_offset, LDS_A_SIZE, BLOCK_N, ROW_STRIDE_B) + + def _lds_elem(rows, row_stride, row, col): + """Element index of (row, col) inside one A- or B-tile of the buffer.""" + if const_expr(lds_layout == "kblock"): + return (col // LOAD_VEC * rows + row) * LOAD_VEC + col % LOAD_VEC + return row * row_stride + col frag_copy_A = fx.make_fragment_like(_pA_s(0)) frag_copy_B = fx.make_fragment_like(_pB_s(0)) @@ -283,8 +370,8 @@ def _load_b_from_lds(rk, buf_offset): col_hi = 16 * rk + 8 for rn in range_constexpr(reg_n): row = wave_n * (reg_n * WMMA_N) + 16 * rn + lane16 - lds_idx_lo = buf_offset + LDS_A_SIZE + row * BLOCK_K_PAD_B + col_lo - lds_idx_hi = buf_offset + LDS_A_SIZE + row * BLOCK_K_PAD_B + col_hi + lds_idx_lo = buf_offset + LDS_A_SIZE + _lds_elem(BLOCK_N, ROW_STRIDE_B, row, col_lo) + lds_idx_hi = buf_offset + LDS_A_SIZE + _lds_elem(BLOCK_N, ROW_STRIDE_B, row, col_hi) v_lo = _v8_load(lds_idx_lo // 8) v_hi = _v8_load(lds_idx_hi // 8) vecs.append(v_lo.shuffle(v_hi, _concat16_mask)) @@ -294,8 +381,8 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset): col_lo = 16 * rk col_hi = 16 * rk + 8 row = wave_m * (reg_m * WMMA_M) + 16 * rm_val + lane16 - lds_idx_lo = buf_offset + row * BLOCK_K_PAD_A + col_lo - lds_idx_hi = buf_offset + row * BLOCK_K_PAD_A + col_hi + lds_idx_lo = buf_offset + _lds_elem(BLOCK_M, ROW_STRIDE_A, row, col_lo) + lds_idx_hi = buf_offset + _lds_elem(BLOCK_M, ROW_STRIDE_A, row, col_hi) v_lo = _v8_load(lds_idx_lo // 8) v_hi = _v8_load(lds_idx_hi // 8) return v_lo.shuffle(v_hi, _concat16_mask) @@ -310,11 +397,30 @@ def _barrier(): has_side_effects=True, ) - def _do_compute_rk(accs_in, rk, buf_offset): + def _do_compute_rk(accs_in, rk, buf_offset, b_vecs): new_accs = list(accs_in) - b_vecs = _load_b_from_lds(rk, buf_offset) + # Each A fragment feeds reg_n back-to-back WMMAs. Emitting its read + # immediately before its first consumer lets the register allocator + # give every rm the same quad, and then the read for rm+1 cannot + # issue until the WMMAs on rm have retired -- the body ends up with + # a full ``lgkmcnt(0)`` drain in front of each group of reg_n WMMAs + # instead of a partial wait. Issuing one fragment ahead is what + # breaks that: the read in flight and the one being consumed need + # different registers, so the wait has something to overlap. + # + # Only survives together with sched_hint. On its own the machine + # scheduler sinks the read back down onto its consumer and the ISA + # comes out unchanged at any prefetch depth; the group barriers are + # what hold the reads up front for the allocator to see. Measured as + # a pair at 256x256x32, worth 5% (109.1 -> 103.9 ms at 16384 cubed) + # and 9 lgkmcnt(0) drains per k-tile down to 4. Neither shows up in + # an instruction histogram -- the counts and the register total are + # the same, only the order changes. + a_next = _load_a_single_from_lds(rk, 0, buf_offset) for rm in range_constexpr(reg_m): - a_vec = _load_a_single_from_lds(rk, rm, buf_offset) + a_vec = a_next + if const_expr(rm + 1 < reg_m): + a_next = _load_a_single_from_lds(rk, rm + 1, buf_offset) for rn in range_constexpr(reg_n): idx = rm * reg_n + rn new_accs[idx] = _wmma_op( @@ -324,6 +430,29 @@ def _do_compute_rk(accs_in, rk, buf_offset): ) return new_accs + def _compute_k_tile(accs_in, buf_offset): + """All reg_k WMMA steps over one LDS buffer. + + Step rk reads B into the registers step rk-1 is still using, so its + reads cannot issue until that step's WMMAs retire, and the wait in + front of them is a full lgkmcnt(0) drain rather than a partial one. + Reading every step's B up front instead does remove that dependence, + and it loses: 2*reg_n*(reg_k-1) more live registers pushes the + allocator into recycling the A fragments harder, and the drains go + from 4 per k-tile to 8. Measured at 256x256x32, 104.0 -> 118.3 ms. + """ + new_accs = list(accs_in) + for rk in range_constexpr(reg_k): + new_accs = _do_compute_rk(new_accs, rk, buf_offset, + _load_b_from_lds(rk, buf_offset)) + return new_accs + + def _sched_k_tile(): + emit = {"vmem": rocdl.sched_vmem, "mfma": rocdl.sched_mfma, + "dsrd": rocdl.sched_dsrd, "dswr": rocdl.sched_dswr} + for group, count in SCHED_PLAN: + emit[group](count) + zero_acc = fx.full(8, 0.0, fx.Float32) accs = [zero_acc for _ in range_constexpr(reg_m * reg_n)] @@ -337,27 +466,33 @@ def _do_compute_rk(accs_in, rk, buf_offset): n_acc = reg_m * reg_n init_state = list(accs) - for iv, state in range(0, num_k_tiles - 1, 1, init=init_state): - s_accs = list(state[:n_acc]) - - read_off = iv % 2 * c_lds_buf_stride - write_off = (1 - iv % 2) * c_lds_buf_stride - - _gmem_load(iv + 1) - - for rk in range_constexpr(reg_k): - s_accs = _do_compute_rk(s_accs, rk, read_off) - + def _one_k_tile(s_accs, read_off, write_off, load_tile): + """Prefetch the next k-tile, consume this one, hand over, barrier.""" + _gmem_load(load_tile) + s_accs = _compute_k_tile(s_accs, read_off) _lds_store(write_off) + if const_expr(sched_hint): + _sched_k_tile() _barrier() - + return s_accs + + # The read buffer alternates with the trip counter, so both offsets are + # values derived from ``iv`` and the body spends ~10 VALU and SALU ops + # per trip on them. Stepping the loop by two fixes each half's parity at + # trace time and folds the arithmetic into ds_load immediates: overhead + # instructions drop from 0.93 to 0.31 per WMMA and VGPRs from 212 to + # 204. It is 1.7% slower (104.0 -> 105.8 ms at 16384 cubed). The loop is + # not issue-bound, so paying more instructions is not what it costs. + for iv, state in range(0, num_k_tiles - 1, 1, init=init_state): + s_accs = list(state[:n_acc]) + s_accs = _one_k_tile(s_accs, iv % 2 * c_lds_buf_stride, + (1 - iv % 2) * c_lds_buf_stride, iv + 1) results = yield list(s_accs) accs = list(results[:n_acc]) last_read_off = ((num_k_tiles - 1) % 2) * c_lds_buf_stride - for rk in range_constexpr(reg_k): - accs = _do_compute_rk(accs, rk, last_read_off) + accs = _compute_k_tile(accs, last_read_off) # ============================================================ # Store results to GMEM through the tiled copy diff --git a/kernels/gemm/rdna3_f16_gemm_autotune.py b/kernels/gemm/rdna3_f16_gemm_autotune.py index fe2336de4..48a706a89 100644 --- a/kernels/gemm/rdna3_f16_gemm_autotune.py +++ b/kernels/gemm/rdna3_f16_gemm_autotune.py @@ -46,12 +46,38 @@ LDS_BYTES = 64 * 1024 # Named by the block tile they produce: (reg_m, reg_n, reg_k, waves_m, waves_n). +TILE_256x256x32 = (4, 4, 2, 4, 4) +TILE_128x256x32 = (4, 4, 2, 2, 4) 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) +# Options a tile needs beyond its shape. Kept next to the ladder rather than in +# the ladder tuples so a Config stays the five tile fields the autotuner sweeps. +# +# 256x256x32 is only reachable unpadded: with K_PAD it wants 80 KB and the pad is +# what the LDS budget cannot afford at that width. Unpadded it is exactly 64 KB. +# The hints are not a tuning knob there but a requirement — without the immediate +# offsets the pad used to give it, LLVM assigns every A fragment the same register +# quad, so each ds_load waits on a full lgkmcnt(0) drain instead of overlapping +# the WMMA stream. Measured at 128x256x32: 53 TFLOP/s without the hints, 84 with. +# +# group_m is the width of the L2 grouping. The wide tiles read a whole 256-row +# band of A per workgroup, so the default 8 walks off the reuse the swizzle is +# there to get; 16 was the best of 1/4/8/16/32 for both, worth 1-2%. +_DEFAULT_OPTS = {"lds_layout": "pad", "sched_hint": False, "group_m": 8} +_TILE_OPTS = { + TILE_256x256x32: {"lds_layout": "kblock", "sched_hint": True, "group_m": 16}, + TILE_128x256x32: {"lds_layout": "pad", "sched_hint": False, "group_m": 16}, +} + + +def tile_opts(tile): + return _TILE_OPTS.get(tuple(tile), _DEFAULT_OPTS) + + # Tile ladder, widest first. # # 128x64x32 exists only on the small-K ladder: with large K a workgroup runs long @@ -62,7 +88,8 @@ # 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_LARGE_K = [TILE_256x256x32, TILE_128x256x32, TILE_128x128x32, TILE_64x64x64, + TILE_32x64x64, TILE_32x32x64] _LADDER_SMALL_K = [TILE_128x128x32, TILE_128x64x32, TILE_64x64x64, TILE_32x64x64, TILE_32x32x64] @@ -80,7 +107,8 @@ def _tile_workgroups(M, N, K, cfg): # 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 + pad = 0 if tile_opts(cfg)["lds_layout"] == "kblock" else K_PAD + if 2 * (block_m + block_n) * (block_k + pad) * 2 > LDS_BYTES: # 2 buffers, 2 bytes/elem return None return (M // block_m) * (N // block_n) @@ -108,8 +136,26 @@ def pick_tile(M, N, K): 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: - + Five exceptions, in order. The first two are the wide tiles, and both are + about operand traffic rather than about filling the machine: a tile reads + ``(BM + BN) / (BM * BN)`` of A and B per output, so 128x256 moves a quarter + less than 128x128 and 256x256 half as much again. Measured against rocBLAS TN + in one thermal window, as a fraction of its time: + + square 4096 8192 12288 16384 20480 + 128x128x32 0.805 0.903 0.919 0.875 0.829 + 128x256x32 0.826 0.933 0.990 0.988 0.928 + 256x256x32 0.797 0.913 0.979 1.012 0.999 + + So the widest tile is not simply best: 256x256 needs roughly 32 workgroups per + CU before its traffic saving outweighs how coarsely its grid lands, and below + that 128x256 leads. Above it the order reverses and stays reversed, because + 128x256 is the one that starts falling off as the footprint grows. + + * 256x256x32 once the grid is worth ~32 workgroups per CU. Only reachable + unpadded; see _TILE_OPTS for why that drags the scheduling hints with it. + * 128x256x32 from ~4 workgroups per CU up. Never behind 128x128x32 anywhere + it applies, by 2-11%. * 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. @@ -124,6 +170,10 @@ def pick_tile(M, N, K): if not feasible: return _ladder_for(K)[0] + if feasible.get(TILE_256x256x32, 0) >= 32 * NUM_CU: + return TILE_256x256x32 + if feasible.get(TILE_128x256x32, 0) >= 4 * NUM_CU: + return TILE_128x256x32 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: @@ -170,6 +220,7 @@ def _build(M, N, K, in_dtype, out_dtype, rounding, reg_m, reg_n, reg_k, waves_m, reg_k=reg_k, waves_m=waves_m, waves_n=waves_n, + **tile_opts((reg_m, reg_n, reg_k, waves_m, waves_n)), ) return launch_fn diff --git a/tests/unit/test_rdna3_gemm_autotune.py b/tests/unit/test_rdna3_gemm_autotune.py index 3988f0cac..9d1ef3b5d 100644 --- a/tests/unit/test_rdna3_gemm_autotune.py +++ b/tests/unit/test_rdna3_gemm_autotune.py @@ -39,6 +39,18 @@ for K in (512, 1024, 2048, 4096) ] +DEEPSEEK_V4_FLASH_DECODE_SHAPES = [ + (M, N, K) + for M in (1, 5, 16) + for N, K in ( + (4096, 1024), + (1024, 4096), + (4096, 512), + (4096, 2048), + (2048, 4096), + ) +] + _shape_id = lambda shape: "x".join(map(str, shape)) # noqa: E731 @@ -87,3 +99,15 @@ def test_feasible_tiles_is_the_ladder_filtered_in_order(shape): assert [tile for tile, _ in reported] == [cfg for cfg in ladder if _tile_workgroups(*shape, cfg) is not None] for tile, workgroups in reported: assert workgroups == _tile_workgroups(*shape, tile) + + +@pytest.mark.parametrize("shape", DEEPSEEK_V4_FLASH_DECODE_SHAPES, ids=_shape_id) +def test_deepseek_v4_flash_decode_shapes_remain_out_of_scope(shape): + """Formal decode rows are small-M tail cases, not tune candidates yet. + + This pins the current scope boundary: these inference shapes have no legal + tile in the existing ladder and would reject at the kernel's ``M % BLOCK_M`` + assertion. Small-M/tail handling is a separate follow-on. + """ + assert feasible_tiles(*shape) == [] + assert all(_tile_workgroups(*shape, tile) is None for tile in _ladder_for(shape[2]))