diff --git a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py index 0aa420b6..cf950626 100644 --- a/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py +++ b/TraceLens/PerfModel/benchmarking/fp4fp6_helpers.py @@ -33,6 +33,16 @@ MX_BLOCK = 32 +def mx_available() -> bool: + """True if any block-scaled MX GEMM path is usable. + Covers the Triton ``tl.dot_scaled`` MXFP4/MXFP6 path as well as aiter's + gfx950 CK ``gemm_a4w4`` MXFP4 path + """ + return ( + triton_available() and (_MXFP4_SUPPORTED or bool(_MXFP6_DTYPE)) + ) or _aiter_mxfp4_ready() + + def triton_available() -> bool: return triton is not None and tl is not None