From 63e6d7f76d81b946c7efe757f3a4d8e523571693 Mon Sep 17 00:00:00 2001 From: Dunfan Lu Date: Tue, 28 Jul 2026 01:39:09 -0700 Subject: [PATCH] Relocate `get_num_xcd` into the dependency-free Triton launcher `get_num_xcd` (AMD CDNA XCD count for `xcd_remap`) is runtime grid logic that generated code invokes as `helion.runtime.get_num_xcd(...)`, so it belongs with the other dependency-free launch helpers. Move it (and `_CUS_PER_XCD`) from `helion/_compat.py` into `helion/runtime/triton/launcher.py`, making the launcher the complete dependency-free home the precompiler bulk-exports. It is re-exported unchanged from `helion.runtime` (so the generated `helion.runtime.get_num_xcd(...)` still resolves); the two importers move to the new location (`autotuner/config_spec.py`, `_testing.py`). `get_num_xcd` stays torch-only, so no import cycle is introduced -- `config_spec` -> `helion.runtime.triton.launcher` is acyclic because the launcher imports no `helion` module. Verified on H100: `import helion` clean, `get_num_xcd` is identical from all re-export sites, `ruff`/`pyrefly` clean (83 baseline), `test_config_api` / `test_runtime` pass. stack-info: PR: https://github.com/pytorch/helion/pull/3180, branch: AmesingFlank/stack/94 --- helion/_compat.py | 41 ------------------------------- helion/_testing.py | 2 +- helion/autotuner/config_spec.py | 2 +- helion/runtime/__init__.py | 2 +- helion/runtime/triton/launcher.py | 41 +++++++++++++++++++++++++++++++ 5 files changed, 44 insertions(+), 44 deletions(-) diff --git a/helion/_compat.py b/helion/_compat.py index 9fdc631db..8056864a8 100644 --- a/helion/_compat.py +++ b/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. diff --git a/helion/_testing.py b/helion/_testing.py index 168fbbf60..a8a5e7d76 100644 --- a/helion/_testing.py +++ b/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( diff --git a/helion/autotuner/config_spec.py b/helion/autotuner/config_spec.py index 33793389c..950771b32 100644 --- a/helion/autotuner/config_spec.py +++ b/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 diff --git a/helion/runtime/__init__.py b/helion/runtime/__init__.py index ef4def9f1..7fee67e32 100644 --- a/helion/runtime/__init__.py +++ b/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: diff --git a/helion/runtime/triton/launcher.py b/helion/runtime/triton/launcher.py index 303ce0bd8..f89423817 100644 --- a/helion/runtime/triton/launcher.py +++ b/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, ...],