fix(perfmodel): sweep tile shapes in the Triton MXFP4/MXFP6 microbenchmark

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
 
@@ -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__)
 
 
@@ -258,13 +271,16 @@ 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
     n = b_nk.shape[0]
     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 +306,42 @@ def _launch_scaled_gemm(
         LHS_DTYPE=lhs_dtype,
         RHS_DTYPE=rhs_dtype,
         PACK=pack,
-        num_warps=8,
+        num_warps=num_warps,
+        **launch_opts,
     )
 
 
+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 +359,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 +553,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
+    )

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,20 +47,17 @@ def _require_cuda_gpu():
 
 def _import_microbench():
     _require_torch()
-
-    return mb
+    return microbench
 
 
 def _import_microbench_rocprof():
     _require_torch()
-
-    return rp
+    return microbench_rocprof
 
 
 def _import_fp4fp6_helpers():
     _require_torch()
-
-    return fp
+    return fp4fp6_helpers
 
 
 class TestMicrobenchUtils:
@@ -248,6 +255,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")