Add FlagGems single-op perf benchmark harness

pytest plugin set + driver script for compiler A/B perf comparison: - reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay - per-shape ttgir dump of actually-used variants with readable naming - cudagraph-based timing with documented fallback semantics Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-16 03:52:07 +00:00
parent 242d74af0e
commit 8ace0d790c
11 changed files with 1562 additions and 1 deletions
@@ -0,0 +1,137 @@
+"""pytest plugin: work around the Hopper fp16 cluster-remote GEMM OOB crash via
+runtime monkeypatch (no upstream source change).
+
+Root cause: fp16 mm uses `_cluster_remote_gemm_kernel` (bf16 doesn't). When
+num_pid_n = cdiv(N,BN) is odd (not a multiple of CLUSTER_SIZE=2) AND the GEMM is
+perfectly tiled (USE_MASK off), the tail cluster's extra CTA reads out of range
+-> cudaErrorIllegalAddress (e.g. N=128256 Llama-3 vocab).
+
+Fix: wrap `cluster_remote_mm_scenario` so only that unsafe combination falls
+back to the general/splitk path; all other fp16 shapes keep the cluster kernel.
+The real backend mm module loads lazily under a synthetic name, so we scan
+sys.modules, force a warmup mm to trigger the load, patch, and re-scan per test.
+"""
+from __future__ import annotations
+
+import sys
+import types
+
+
+def _patch_loaded() -> list:
+    import triton
+
+    patched = []
+    for nm, m in list(sys.modules.items()):
+        if not isinstance(m, types.ModuleType):
+            continue
+        fn = getattr(m, "cluster_remote_mm_scenario", None)
+        if not isinstance(fn, types.FunctionType):
+            continue
+        if getattr(fn, "_mm_cluster_guarded", False):
+            continue
+        bn = getattr(m, "TLE_REMOTE_BN", 256)
+        bm = getattr(m, "TLE_REMOTE_BM", 64)
+        bk = getattr(m, "TLE_REMOTE_BK", 64)
+        cs = getattr(m, "TLE_CLUSTER_SIZE", 2)
+
+        def make(orig, bm, bn, bk, cs):
+            def wrapped(a, b, c, M, N, K):
+                if not orig(a, b, c, M, N, K):
+                    return False
+                use_mask = (M % bm != 0) or (N % bn != 0) or (K % bk != 0)
+                # odd N-tile count + no mask -> tail cluster over-runs N
+                if (triton.cdiv(N, bn) % cs != 0) and not use_mask:
+                    return False
+                return True
+            wrapped._mm_cluster_guarded = True
+            return wrapped
+
+        m.cluster_remote_mm_scenario = make(fn, bm, bn, bk, cs)
+        patched.append(nm)
+    return patched
+
+
+def _guarded_modules() -> list:
+    """Loaded modules whose cluster_remote_mm_scenario is already our wrapper
+    (cumulative state, unlike _patch_loaded()'s newly-wrapped-only return)."""
+    out = []
+    for nm, m in list(sys.modules.items()):
+        if not isinstance(m, types.ModuleType):
+            continue
+        fn = getattr(m, "cluster_remote_mm_scenario", None)
+        if isinstance(fn, types.FunctionType) and getattr(fn, "_mm_cluster_guarded", False):
+            out.append(nm)
+    return out
+
+
+def _is_hopper():
+    """True iff a CUDA device with capability >= 9 (the only HW that hits the
+    cluster-remote path) is visible. Returns False on any error / no CUDA."""
+    try:
+        import torch
+        if not torch.cuda.is_available():
+            return False
+        return torch.cuda.get_device_capability()[0] >= 9
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _warmup_load_backend():
+    """Trigger the lazy backend mm module load via one tiny 64x64 fp16 mm (too
+    small for the cluster path, so it can't trip the bug). A CUDA error here
+    means a pre-broken context -> surface loudly; only import/attr errors are
+    swallowed."""
+    try:
+        import torch
+        import flag_gems
+        if not torch.cuda.is_available():
+            return
+        a = torch.randn(64, 64, dtype=torch.float16, device="cuda")
+        b = torch.randn(64, 64, dtype=torch.float16, device="cuda")
+        with flag_gems.use_gems():
+            torch.mm(a, b)
+            torch.cuda.synchronize()
+    except (ImportError, AttributeError) as e:
+        print(f"[mm-cluster-fix-plugin] warmup skipped (benign): {e!r}",
+              file=sys.stderr, flush=True)
+    except Exception as e:  # noqa: BLE001
+        # e.g. a CUDA error — do NOT hide it; the harness needs to see it.
+        print(f"[mm-cluster-fix-plugin] WARNING: warmup mm failed unexpectedly: {e!r}",
+              file=sys.stderr, flush=True)
+
+
+def pytest_configure(config):
+    _patch_loaded()
+    _warmup_load_backend()
+    _patch_loaded()
+    guarded = _guarded_modules()
+    if guarded:
+        print(
+            "[mm-cluster-fix-plugin] guarded cluster_remote_mm_scenario "
+            f"(odd N-tile unmasked OOB) in: {guarded}",
+            file=sys.stderr,
+            flush=True,
+        )
+    elif _is_hopper():
+        # On Hopper with nothing patched, the guard is a silent no-op and fp16 mm
+        # will crash; fail loudly rather than risk an illegal-access.
+        raise RuntimeError(
+            "[mm-cluster-fix-plugin] FATAL: on a Hopper (cap>=9) device but "
+            "cluster_remote_mm_scenario was not found in any loaded module after "
+            "warmup. The fp16 mm OOB guard is NOT active; aborting rather than "
+            "risk an illegal-memory-access crash. (FlagGems mm module layout may "
+            "have changed.)"
+        )
+    else:
+        # Non-Hopper: the cluster path never runs, so an unpatched state is fine.
+        print(
+            "[mm-cluster-fix-plugin] no cluster_remote_mm_scenario found; "
+            "non-Hopper device, guard not needed (no-op).",
+            file=sys.stderr,
+            flush=True,
+        )
+
+
+def pytest_runtest_setup(item):
+    # backend module may load lazily on the first gems call; re-scan (idempotent).
+    _patch_loaded()