zl_bench/_mm_cluster_fix_plugin.py

"""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:
        from _term_style import tag
        print(f"{tag('[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.
        from _term_style import tag
        print(f"{tag('[mm-cluster-fix-plugin]')} WARNING: warmup mm failed "
              f"unexpectedly: {e!r}",
              file=sys.stderr, flush=True)


def pytest_configure(config):
    from _term_style import tag
    _patch_loaded()
    _warmup_load_backend()
    _patch_loaded()
    guarded = _guarded_modules()
    if guarded:
        print(
            f"{tag('[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(
            f"{tag('[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()