From 2f1672b407ada05b29fc663ce57627dbee836e99 Mon Sep 17 00:00:00 2001 From: Mehdi Saeedi <32145191+mehdi-saeedi@users.noreply.github.com> Date: Wed, 7 Oct 2026 13:51:41 -0400 Subject: [PATCH 1/2] fix(perfmodel): sweep tile shapes in the Triton MXFP4/MXFP6 microbenchmark The tl.dot_scaled MXFP4/MXFP6 GEMM benchmark launched one fixed tile (128x128x256, 8 warps), which suits gfx950 but is a poor fit for RDNA4 (gfx1201), where smaller BLOCK_K tiles are much faster on the same shapes. The reported peak therefore depended on whether that one tile happened to fit the target. Try a small set of tiles per shape (MX_TILE_CONFIGS, starting with the original tile so no target can regress) and keep the fastest; tiles that do not divide the shape or fail to compile are skipped. Each entry also carries num_stages: None keeps Triton's backend default (the original behaviour), and a single-stage tile is included because dropping software pipelining helps further on RDNA4. Also import the modules in the test helpers _import_microbench, _import_microbench_rocprof and _import_fp4fp6_helpers, which returned undefined names, so every torch-enabled test using them failed with a NameError (CI without torch skips them). Co-authored-by: Cursor --- .../PerfModel/benchmarking/fp4fp6_helpers.py | 109 ++++++++++++++++-- tests/test_perfmodel_benchmarking.py | 44 +++++++ 2 files changed, 143 insertions(+), 10 deletions(-) diff --git a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py index 985c7b076..c8eac18f6 100644 --- a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py +++ b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py @@ -15,7 +15,7 @@ from __future__ import annotations import logging -from typing import Tuple +from typing import Optional, Tuple import torch @@ -258,6 +258,8 @@ def _launch_scaled_gemm( block_m: int = 128, block_n: int = 128, block_k: int = 256, + num_warps: int = 8, + num_stages: Optional[int] = None, ) -> None: m, packed_k = a.shape k = packed_k * pack @@ -265,6 +267,7 @@ def _launch_scaled_gemm( w = b_nk.T # kernel expects (K // PACK, N) grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n)) + launch_opts = {} if num_stages is None else {"num_stages": num_stages} _mx_dot_scaled_kernel[grid]( a, w, @@ -290,10 +293,56 @@ def _launch_scaled_gemm( LHS_DTYPE=lhs_dtype, RHS_DTYPE=rhs_dtype, PACK=pack, - num_warps=8, + num_warps=num_warps, + **launch_opts, ) +# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) tried for every shape; the +# fastest is kept. num_stages=None leaves Triton's backend default. The first +# entry is the original fixed tile, which suits gfx950; on RDNA4 the smaller +# BLOCK_K tiles are much faster, and single-stage (no software pipelining) +# helps further. +MX_TILE_CONFIGS: Tuple[Tuple[int, int, int, int, Optional[int]], ...] = ( + (128, 128, 256, 8, None), + (128, 128, 128, 8, None), + (128, 128, 64, 8, None), + (64, 128, 64, 4, None), + (128, 128, 64, 4, 1), +) + + +def _bench_best_tile(m: int, n: int, k: int, run, *, warmup, rep, do_bench_fn): + """ + Best TFLOPS of ``run(block_m, block_n, block_k, num_warps, num_stages)`` over the + ``MX_TILE_CONFIGS`` that evenly divide the shape (the kernel has no + bounds masks). Tiles that fail to compile or launch are skipped; if every + applicable tile fails, the last error is raised. + """ + best, error = 0.0, None + for block_m, block_n, block_k, num_warps, num_stages in MX_TILE_CONFIGS: + if m % block_m or n % block_n or k % block_k: + continue + try: + ms = do_bench_fn( + lambda: run(block_m, block_n, block_k, num_warps, num_stages), + warmup=warmup, + rep=rep, + ) + except Exception as e: + logger.debug( + "MX tile %s failed", + (block_m, block_n, block_k, num_warps, num_stages), + exc_info=True, + ) + error = e + continue + best = max(best, gemm_tflops(m, n, k, ms)) + if best == 0.0 and error is not None: + raise error + return best + + def bench_mxfp4_gemm( m: int, n: int, @@ -311,11 +360,32 @@ def bench_mxfp4_gemm( dev = torch.device(f"cuda:{device}") a, b, sa, sb, c = prepare_mxfp4_gemm(m, n, k, dev) - def _run() -> None: - _launch_scaled_gemm(a, b, sa, sb, c, "e2m1", "e2m1", pack=2) + def _run( + block_m: int, + block_n: int, + block_k: int, + num_warps: int, + num_stages: Optional[int], + ) -> None: + _launch_scaled_gemm( + a, + b, + sa, + sb, + c, + "e2m1", + "e2m1", + pack=2, + block_m=block_m, + block_n=block_n, + block_k=block_k, + num_warps=num_warps, + num_stages=num_stages, + ) - ms = do_bench_fn(_run, warmup=warmup, rep=rep) - return gemm_tflops(m, n, k, ms) + return _bench_best_tile( + m, n, k, _run, warmup=warmup, rep=rep, do_bench_fn=do_bench_fn + ) _AITER_MXFP4_CACHE: dict = { @@ -484,10 +554,29 @@ def bench_mxfp6_gemm( else: a, b, sa, sb, c = prepare_mxfp6_gemm(m, n, k, dev) - def _run() -> None: + def _run( + block_m: int, + block_n: int, + block_k: int, + num_warps: int, + num_stages: Optional[int], + ) -> None: _launch_scaled_gemm( - a, b, sa, sb, c, _MXFP6_DTYPE, _MXFP6_DTYPE, pack=_MXFP6_PACK + a, + b, + sa, + sb, + c, + _MXFP6_DTYPE, + _MXFP6_DTYPE, + pack=_MXFP6_PACK, + block_m=block_m, + block_n=block_n, + block_k=block_k, + num_warps=num_warps, + num_stages=num_stages, ) - ms = do_bench_fn(_run, warmup=warmup, rep=rep) - return gemm_tflops(m, n, k, ms) + return _bench_best_tile( + m, n, k, _run, warmup=warmup, rep=rep, do_bench_fn=do_bench_fn + ) diff --git a/tests/test_perfmodel_benchmarking.py b/tests/test_perfmodel_benchmarking.py index 6322a5aac..0f3e05cc0 100644 --- a/tests/test_perfmodel_benchmarking.py +++ b/tests/test_perfmodel_benchmarking.py @@ -37,18 +37,21 @@ def _require_cuda_gpu(): def _import_microbench(): _require_torch() + from TraceLens.PerfModel.benchmarking import microbench as mb return mb def _import_microbench_rocprof(): _require_torch() + from TraceLens.PerfModel.benchmarking import microbench_rocprof as rp return rp def _import_fp4fp6_helpers(): _require_torch() + from TraceLens.PerfModel.benchmarking import fp4fp6_helpers as fp return fp @@ -248,6 +251,47 @@ def fake_bench(_fn, *, warmup, rep): == 0.0 ) + def test_mx_tile_sweep_keeps_original_tile_first(self): + fp = _import_fp4fp6_helpers() + assert fp.MX_TILE_CONFIGS[0] == (128, 128, 256, 8, None) + + def test_best_tile_keeps_fastest_and_skips_non_dividing(self): + fp = _import_fp4fp6_helpers() + ms_by_tile = { + (128, 128, 256, 8, None): 8.0, + (128, 128, 128, 8, None): 2.0, + (128, 128, 64, 8, None): 3.0, + (64, 128, 64, 4, None): 4.0, + (128, 128, 64, 4, 1): 1.0, + } + tried = [] + + def run(*tile): + tried.append(tile) + + def fake_bench(fn, *, warmup, rep): + fn() + return ms_by_tile[tried[-1]] + + m, n, k = 256, 256, 1152 # 1152 is not a multiple of BLOCK_K=256 + best = fp._bench_best_tile( + m, n, k, run, warmup=1, rep=1, do_bench_fn=fake_bench + ) + assert (128, 128, 256, 8, None) not in tried + assert (128, 128, 64, 4, 1) in tried + assert best == pytest.approx(fp.gemm_tflops(m, n, k, 1.0)) + + def test_best_tile_raises_when_every_tile_fails(self): + fp = _import_fp4fp6_helpers() + + def fake_bench(_fn, *, warmup, rep): + raise RuntimeError("out of resources") + + with pytest.raises(RuntimeError, match="out of resources"): + fp._bench_best_tile( + 256, 256, 256, lambda *t: None, warmup=1, rep=1, do_bench_fn=fake_bench + ) + @pytest.mark.gpu @pytest.mark.skipif(not HAS_TORCH, reason="torch not installed") From 40738fbfe94a8119ea469c0e1fc00eb2a7924a32 Mon Sep 17 00:00:00 2001 From: Mehdi Saeedi <32145191+mehdi-saeedi@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:46:28 -0400 Subject: [PATCH 2/2] Address review: gather the benchmarking imports and the tile list at the top - The test module imports fp4fp6_helpers, microbench and microbench_rocprof once at the top, behind an ImportError guard because they import torch and CI runs without it; the _import_* helpers keep their torch skip and return the module. - MX_TILE_CONFIGS moves up beside MX_BLOCK with the module's other constants. Co-authored-by: Cursor --- .../PerfModel/benchmarking/fp4fp6_helpers.py | 27 +++++++++---------- tests/test_perfmodel_benchmarking.py | 22 ++++++++------- 2 files changed, 26 insertions(+), 23 deletions(-) diff --git a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py index c8eac18f6..d436a7546 100644 --- a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py +++ b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py @@ -35,6 +35,19 @@ # https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf MX_BLOCK = 32 +# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) tried for every shape; the +# fastest is kept. num_stages=None leaves Triton's backend default. The first +# entry is the original fixed tile, which suits gfx950; on RDNA4 the smaller +# BLOCK_K tiles are much faster, and single-stage (no software pipelining) +# helps further. +MX_TILE_CONFIGS: Tuple[Tuple[int, int, int, int, Optional[int]], ...] = ( + (128, 128, 256, 8, None), + (128, 128, 128, 8, None), + (128, 128, 64, 8, None), + (64, 128, 64, 4, None), + (128, 128, 64, 4, 1), +) + logger = logging.getLogger(__name__) @@ -298,20 +311,6 @@ def _launch_scaled_gemm( ) -# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) tried for every shape; the -# fastest is kept. num_stages=None leaves Triton's backend default. The first -# entry is the original fixed tile, which suits gfx950; on RDNA4 the smaller -# BLOCK_K tiles are much faster, and single-stage (no software pipelining) -# helps further. -MX_TILE_CONFIGS: Tuple[Tuple[int, int, int, int, Optional[int]], ...] = ( - (128, 128, 256, 8, None), - (128, 128, 128, 8, None), - (128, 128, 64, 8, None), - (64, 128, 64, 4, None), - (128, 128, 64, 4, 1), -) - - def _bench_best_tile(m: int, n: int, k: int, run, *, warmup, rep, do_bench_fn): """ Best TFLOPS of ``run(block_m, block_n, block_k, num_warps, num_stages)`` over the diff --git a/tests/test_perfmodel_benchmarking.py b/tests/test_perfmodel_benchmarking.py index 0f3e05cc0..e2cdb29c9 100644 --- a/tests/test_perfmodel_benchmarking.py +++ b/tests/test_perfmodel_benchmarking.py @@ -21,6 +21,16 @@ resolve_physical_device, ) +# These modules import torch; without it they stay None and the tests that use them skip in _require_torch. +try: + from TraceLens.PerfModel.benchmarking import ( + fp4fp6_helpers, + microbench, + microbench_rocprof, + ) +except ImportError: + fp4fp6_helpers = microbench = microbench_rocprof = None + HAS_TORCH = importlib.util.find_spec("torch") is not None @@ -37,23 +47,17 @@ def _require_cuda_gpu(): def _import_microbench(): _require_torch() - from TraceLens.PerfModel.benchmarking import microbench as mb - - return mb + return microbench def _import_microbench_rocprof(): _require_torch() - from TraceLens.PerfModel.benchmarking import microbench_rocprof as rp - - return rp + return microbench_rocprof def _import_fp4fp6_helpers(): _require_torch() - from TraceLens.PerfModel.benchmarking import fp4fp6_helpers as fp - - return fp + return fp4fp6_helpers class TestMicrobenchUtils: