From 7370720e146ca6898e3aaebad77dcc804154c1e6 Mon Sep 17 00:00:00 2001 From: xudoyuan Date: Fri, 12 Jun 2026 08:58:15 +0000 Subject: [PATCH] Add SUPPORTED_ARCHS table and is_arch_supported() helper Provide a single source of truth for which GPU architectures FlyDSL ships kernels for. Downstream consumers (e.g. aiter) can gate availability on is_arch_supported() so that importing FlyDSL kernels on an unsupported arch (e.g. gfx1100/RDNA3) is skipped cleanly instead of crashing during config registration with a KeyError on SMEM_CAPACITY_MAP. Co-Authored-By: Claude Opus 4 --- python/flydsl/runtime/device.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/python/flydsl/runtime/device.py b/python/flydsl/runtime/device.py index ed5fd47b2..ac95fc91c 100644 --- a/python/flydsl/runtime/device.py +++ b/python/flydsl/runtime/device.py @@ -91,3 +91,28 @@ def is_rdna_arch(arch: Optional[str] = None) -> bool: if arch.startswith("gfx120"): return True return False + + +# Architectures FlyDSL ships kernels for. Single source of truth for +# "does FlyDSL run here" — keep in sync with the kernels under kernels/. +# Downstream consumers (e.g. aiter) gate availability on this so that +# importing FlyDSL kernels on an unsupported arch is skipped cleanly +# instead of crashing during config registration. +SUPPORTED_ARCHS = frozenset( + { + "gfx942", # MI300A / MI300X (CDNA3) + "gfx950", # MI350 (CDNA4) + "gfx1151", # RDNA3.5 (Strix Halo) + "gfx1201", # RDNA4 + "gfx1250", # MI450 + } +) + + +def is_arch_supported(arch: Optional[str] = None) -> bool: + """Whether FlyDSL ships kernels for *arch* (current GPU if None).""" + if arch is None: + arch = get_rocm_arch() + if not arch: + return False + return arch.lower() in SUPPORTED_ARCHS