Relocate get_num_xcd into the dependency-free Triton launcher

helion/_compat.py

@@ -480,47 +480,6 @@ def supports_amd_cdna_tunables() -> bool:
         return False
 
 
-# CUs per XCD by base CDNA architecture.  Used to derive the live,
-# partition-visible XCD count from the observed CU count (see get_num_xcd).
-_CUS_PER_XCD: dict[str, int] = {
-    "gfx942": 38,  # CDNA3 (MI300)
-    "gfx950": 32,  # CDNA4 (MI350)
-    "gfx951": 32,  # CDNA4 (MI355)
-}
-
-
-def get_num_xcd(device: torch.device | int | None = None) -> int:
-    """Number of XCDs visible for ``device`` on AMD CDNA, else ``1``.
-
-    Derived from the live, partition-visible compute-unit count rather than the
-    architecture name, so MI300A (6 XCDs) and compute-partition modes such as CPX
-    (which expose a single XCD) are handled correctly.  Returns ``1`` -- which
-    disables xcd_remap -- for unknown architectures or a CU count that does not
-    look like an integer number of XCDs.
-    """
-    if not torch.cuda.is_available():
-        return 1
-    try:
-        props = torch.cuda.get_device_properties(
-            device if device is not None else torch.cuda.current_device()
-        )
-    except Exception:
-        return 1
-    arch = getattr(props, "gcnArchName", None)
-    if not arch:
-        return 1
-    cus_per_xcd = _CUS_PER_XCD.get(arch.split(":")[0])
-    if cus_per_xcd is None:
-        return 1
-    cu_count = props.multi_processor_count
-    num_xcd = round(cu_count / cus_per_xcd)
-    # Tolerate harvested parts, but bail out (return 1) if the live CU count does
-    # not look like an integer number of XCDs.
-    if num_xcd < 1 or abs(num_xcd * cus_per_xcd - cu_count) > cus_per_xcd // 4:
-        return 1
-    return num_xcd
-
-
 def device_num_sm(device: torch.device | int | None = None) -> int:
     """SM/CU count for ``device`` (or the current device), or 1 if unavailable.
 

helion/_testing.py

@@ -356,8 +356,8 @@ def skipUnlessMultiXCD(reason: str) -> Callable[[Callable], Callable]:
     Single-XCD parts and CPX-partitioned devices (which expose one XCD) are
     skipped, since xcd_remap is a no-op there.
     """
-    from helion._compat import get_num_xcd
     from helion._compat import supports_amd_cdna_tunables
+    from helion.runtime.triton.launcher import get_num_xcd
 
     # Defers check to test execution time to avoid CUDA init during pytest-xdist collection.
     return skipIfFn(

helion/autotuner/config_spec.py

@@ -20,7 +20,6 @@
 
 from .._compat import _regs_per_block
 from .._compat import device_num_sm
-from .._compat import get_num_xcd
 from .._compat import num_compute_units
 from .._compat import supports_amd_cdna_tunables
 from .._compat import supports_maxnreg
@@ -60,6 +59,7 @@
 from .._compiler.cute.tcgen05_config import Tcgen05ClusterM2SearchConstraints
 from .._compiler.cute.tcgen05_constants import TCGEN05_TWO_CTA_MAX_K_TILES
 from ..exc import InvalidConfig
+from ..runtime.triton.launcher import get_num_xcd
 from .block_id_sequence import BlockIdSequence
 from .block_id_sequence import _BlockIdItem
 from .block_id_sequence import _PowerOfTwoBlockIdItem

helion/runtime/__init__.py

@@ -20,7 +20,6 @@
 
 from .. import _compat as _compat  # ensure Triton compatibility patches run
 from .. import exc
-from .._compat import get_num_xcd as get_num_xcd
 from .._compiler.cute.strategies import tcgen05_default_epilogue_tile_expr
 from .._compiler.cute.strategies import tcgen05_explicit_d_store_tile_expr
 from .._compiler.cute.strategies import tcgen05_smem_layout_expr
@@ -30,6 +29,7 @@
 from .settings import is_pallas_interpret as _module_is_pallas_interpret
 from .triton.launcher import default_launcher as _triton_default_launcher
 from .triton.launcher import get_num_sm as _triton_get_num_sm
+from .triton.launcher import get_num_xcd as get_num_xcd
 from .triton.launcher import set_triton_allocator as set_triton_allocator
 
 if TYPE_CHECKING:

helion/runtime/triton/launcher.py

@@ -112,6 +112,47 @@ def get_num_sm(device: torch.device, *, reserved_sms: int = 0) -> int:
     return max(available_sms - reserved_sms, 1)
 
 
+# CUs per XCD by base CDNA architecture.  Used to derive the live,
+# partition-visible XCD count from the observed CU count (see get_num_xcd).
+_CUS_PER_XCD: dict[str, int] = {
+    "gfx942": 38,  # CDNA3 (MI300)
+    "gfx950": 32,  # CDNA4 (MI350)
+    "gfx951": 32,  # CDNA4 (MI355)
+}
+
+
+def get_num_xcd(device: torch.device | int | None = None) -> int:
+    """Number of XCDs visible for ``device`` on AMD CDNA, else ``1``.
+
+    Derived from the live, partition-visible compute-unit count rather than the
+    architecture name, so MI300A (6 XCDs) and compute-partition modes such as CPX
+    (which expose a single XCD) are handled correctly.  Returns ``1`` -- which
+    disables xcd_remap -- for unknown architectures or a CU count that does not
+    look like an integer number of XCDs.
+    """
+    if not torch.cuda.is_available():
+        return 1
+    try:
+        props = torch.cuda.get_device_properties(
+            device if device is not None else torch.cuda.current_device()
+        )
+    except Exception:
+        return 1
+    arch = getattr(props, "gcnArchName", None)
+    if not arch:
+        return 1
+    cus_per_xcd = _CUS_PER_XCD.get(arch.split(":")[0])
+    if cus_per_xcd is None:
+        return 1
+    cu_count = props.multi_processor_count
+    num_xcd = round(cu_count / cus_per_xcd)
+    # Tolerate harvested parts, but bail out (return 1) if the live CU count does
+    # not look like an integer number of XCDs.
+    if num_xcd < 1 or abs(num_xcd * cus_per_xcd - cu_count) > cus_per_xcd // 4:
+        return 1
+    return num_xcd
+
+
 def default_launcher(
     triton_kernel: object,
     grid: tuple[int, ...],