[Kernel][gfx1250] Fix FlyDSL MXScale import and arch gates

Author: aoli26

aiter/__init__.py

@@ -127,7 +127,7 @@ def getLogger():
         gemm_mxfp8,
         gemm_mxa8w4,
     )
-except ImportError:
+except (ImportError, RuntimeError, OSError, KeyError):
     pass
 
 # Import Triton-based communication primitives from ops.triton.comms (optional, only if Iris is available)

aiter/ops/flydsl/mxscale_gemm.py

@@ -23,7 +23,7 @@
 from torch import Tensor
 
 from aiter import logger
-from aiter.jit.utils.chip_info import get_gfx
+from aiter.jit.utils.chip_info import get_gfx_runtime
 
 from .mxscale_layout import (
     SCALE_BLOCK,
@@ -273,7 +273,7 @@ def flydsl_mxscale_gemm(
         raise ValueError(
             f"data_format must be one of {_VALID_FORMATS}, got {data_format!r}"
         )
-    cur_gfx = get_gfx()
+    cur_gfx = get_gfx_runtime()
     if cur_gfx != _TARGET_GFX:
         raise RuntimeError(
             f"flydsl_mxscale_gemm requires {_TARGET_GFX}, current arch is {cur_gfx!r}"

aiter/ops/flydsl/test_flydsl_mxscale_gemm.py

@@ -334,6 +334,20 @@ def test_format_named_wrappers_reject_data_format_kwarg():
         gemm_mxa8w4(a, b, s, s, data_format="a8w4")
 
 
+def test_runtime_arch_gate_ignores_gpu_archs(monkeypatch):
+    from aiter.jit.utils import chip_info
+
+    monkeypatch.setenv("GPU_ARCHS", "gfx950")
+    chip_info.get_gfx.cache_clear()
+    chip_info.get_gfx_custom_op_core.cache_clear()
+
+    a = torch.zeros((32,), dtype=torch.uint8)
+    b = torch.zeros((32,), dtype=torch.uint8)
+    s = torch.zeros((1, 1), dtype=torch.uint8)
+    with pytest.raises(ValueError, match="A and B must be 2-D"):
+        flydsl_mxscale_gemm(a, b, s, s, data_format="fp8")
+
+
 # ---------------------------------------------------------------------------
 # GPU correctness tests (gfx1250 + flydsl required, gated above)
 # ---------------------------------------------------------------------------