From 0ec49b032aee18b703072375851aa59aa64a5164 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 10:00:50 -0400 Subject: [PATCH 01/16] fix(rocm): resolve the build architecture in one place, not three Three call sites independently answered "what architecture are we building for", and on CDNA4 they disagreed. Measured on a gfx950 host with FLASHINFER_ROCM_ARCH_LIST unset, before this change: ACTUAL device arch gfx950 validate_flashinfer_rocm_arch(arch_list=None) ['gfx942'] <-- wrong CompilationContext().TARGET_ROCM_ARCHS ['gfx950'] resolve_aiter_build_arch() gfx950 So the JIT validated gfx942 while compiling for gfx950. The check exists to catch "your PyTorch was not built for this architecture", and it was asking about an architecture nobody was building for: vacuous on a PyTorch carrying both, and a spurious hard failure on an arch-specific build carrying only gfx950. The cause was `os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942")` reached from two functions in hip_utils, plus two more `return "gfx942"` lines in CompilationContext._auto_detect_archs. Add hip_utils.resolve_target_archs() -- explicit argument, then the env var, then the architectures actually present, then every supported architecture with a warning -- and route validate_rocm_arch, validate_flashinfer_rocm_arch, CompilationContext and aot_hip through it. _auto_detect_archs goes away; it was private and had no other caller. The last-resort fallback changes from "gfx942" to every supported architecture. On a GPU-less build host the old literal was not conservative, it was a guess that silently produced a gfx942-only artifact; a fat build is slower but correct wherever it lands, and the warning says how to make it cheap again. Detection uses rocminfo rather than torch.cuda, so the resolver adds no torch dependency to a module that must stay importable without one. test_defaults_to_gfx942_when_no_env_and_no_arg asserted the wrong constant -- which is how the constant survived -- and now pins the detected architecture. Measured after, same host and script: all four rows report gfx950. Also verified FLASHINFER_ROCM_ARCH_LIST=gfx942 is still honoured on a gfx950 box (cross-compiling stays possible), and that a container with no /dev/kfd resolves to "gfx942,gfx950" with the warning and without importing torch. 169 passed across test_hip_utils.py, test_aiter_build_arch_hip.py and test_arch_caps_hip.py on gfx950. Co-Authored-By: Claude --- flashinfer/aot_hip.py | 36 ++++----- flashinfer/compilation_context_hip.py | 35 ++------- flashinfer/hip_utils.py | 69 +++++++++++++++-- tests/rocm_tests/test_hip_utils.py | 103 +++++++++++++++++++++++++- 4 files changed, 187 insertions(+), 56 deletions(-) diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 5912977f27..5c4ee8f18c 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -208,25 +208,27 @@ def compile_and_package_modules( final_config.update(config) config = final_config - # ROCm Arch: Ensure env var is set or create/validate using CompilationContext + # ROCm arch: resolve once, then validate. + # + # Publishing the result back into the environment is deliberate, not + # incidental bookkeeping: the AITER shim resolves its own build architecture + # from FLASHINFER_ROCM_ARCH_LIST (jit/aiter_source.py), and an AOT build has + # no other channel to tell it what this build targets. Without this, a shim + # built during an AOT run on a mixed or GPU-less host can disagree with the + # kernels it is packaged alongside. + # + # It is a process-global side effect that outlives the call, which is worth + # replacing with an explicit parameter threaded through the AOT -> JIT + # boundary. That is a wider change than this one; leaving the lifetime + # unchanged here keeps this commit to the resolution bug it is fixing. from .compilation_context_hip import CompilationContext + from .hip_utils import resolve_target_archs - if "FLASHINFER_ROCM_ARCH_LIST" not in os.environ: - # Auto-detect or use default by creating a local context - compilation_context = CompilationContext() - detected_archs = ",".join(sorted(compilation_context.TARGET_ROCM_ARCHS)) - os.environ["FLASHINFER_ROCM_ARCH_LIST"] = detected_archs - if verbose: - print(f"Auto-detected ROCm architectures: {detected_archs}") - else: - # Validate provided arch list by creating a local context - arch_list = os.environ["FLASHINFER_ROCM_ARCH_LIST"] - CompilationContext() # Validates arch_list set via env var - if verbose: - print(f"Using ROCm architectures: {arch_list}") - - # Verify paths are correct - rocm_arch_list = os.environ["FLASHINFER_ROCM_ARCH_LIST"] + rocm_arch_list = resolve_target_archs() + os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list + CompilationContext() # validates the resolved list, raising on a bad one + if verbose: + print(f"Target ROCm architectures: {rocm_arch_list}") # Print summary if verbose: diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index 28eb694f76..ee8cdd74f0 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -18,12 +18,9 @@ """ import logging -import os -import torch from . import hip_utils -from .arch_caps import normalize_arch logger = logging.getLogger(__name__) @@ -50,12 +47,11 @@ def __init__(self): """ import torch.utils.cpp_extension as torch_cpp_ext - # Get architecture list from env or auto-detect - arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") - if arch_list is None: - arch_list = self._auto_detect_archs() - if arch_list: - logger.info(f"Auto-detected ROCm architectures: {arch_list}") + # One resolver for every path that asks "what are we building for", so + # this cannot disagree with the validation in hip_utils -- it used to, + # returning gfx950 here while validation checked gfx942. + arch_list = hip_utils.resolve_target_archs() + logger.info(f"Target ROCm architectures: {arch_list}") # Comprehensive validation (all 3 checks) self.arch_flags, self.TARGET_ROCM_ARCHS = ( @@ -64,27 +60,6 @@ def __init__(self): ) ) - def _auto_detect_archs(self) -> str: - """Auto-detect ROCm architectures from supported system devices. - - Only devices whose gcnArchName is in FLASHINFER_SUPPORTED_ROCM_ARCHS are - considered, so unsupported integrated GPUs are silently ignored here rather - than being passed on to validate_flashinfer_rocm_arch for filtering. - """ - try: - indices = hip_utils.get_supported_device_indices() - if indices: - archs = { - normalize_arch(torch.cuda.get_device_properties(i).gcnArchName) - for i in indices - } - return ",".join(sorted(archs)) - logger.warning("No supported ROCm devices detected, defaulting to gfx942") - return "gfx942" - except Exception as e: - logger.warning(f"Failed to auto-detect ROCm device architectures: {e}") - return "gfx942" - def get_hipcc_flags_list(self) -> list[str]: """ Generate hipcc compiler flags for target architectures. diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 6a91f0b76b..9016dcb49a 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -3,16 +3,77 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import logging # arch_caps imports nothing (in particular, not torch), so importing it here # keeps this module safe to import before the HIP runtime starts -- which # tests/conftest.py relies on to set HIP_VISIBLE_DEVICES first. from .arch_caps import normalize_arch +logger = logging.getLogger(__name__) + # AMDGPU archs supported by amd-flashinfer FLASHINFER_SUPPORTED_ROCM_ARCHS = ["gfx942", "gfx950"] +def resolve_target_archs(arch_list: str = None) -> str: + """Return the architectures to build for, as a comma-separated string. + + The single answer to "what are we compiling for". Resolution order: + + 1. ``arch_list``, when a caller passes one explicitly. + 2. ``FLASHINFER_ROCM_ARCH_LIST``. + 3. The architectures of the supported GPUs actually present. + 4. Every architecture FlashInfer supports, with a warning. + + Step 4 replaces a hard-coded ``"gfx942"`` that three call sites reached + independently. On a CDNA4 host that literal was not a conservative default + but a wrong answer: ``validate_flashinfer_rocm_arch(arch_list=None)`` + returned ``{"gfx942"}`` on a gfx950 device while ``CompilationContext`` + compiled for gfx950, so the check that exists to catch "your PyTorch was not + built for this architecture" was validating an architecture nobody was + building for. Vacuous on a PyTorch carrying both; a spurious hard failure on + an arch-specific build that carries only gfx950. + + Building for everything we support is the honest fallback when we cannot + tell: slower and fatter, but correct on whichever card the result lands on, + which a guess is not. It is warned so a GPU-less build host is told to set + the variable rather than silently paying for it. + + Detection uses rocminfo rather than ``torch.cuda`` so this stays callable + before the HIP runtime starts, and keeps this module importable without + torch -- see ``tests/rocm_tests/test_arch_caps_hip.py``. + """ + import os + + if arch_list: + return arch_list + + from_env = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") + if from_env: + return from_env + + detected = sorted( + { + arch + for arch, _ in rocminfo_gpu_agents() + if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS + } + ) + if detected: + return ",".join(detected) + + fallback = ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + logger.warning( + "No supported AMD GPU detected and FLASHINFER_ROCM_ARCH_LIST is unset; " + "building for every supported architecture (%s). This is slower than " + "targeting one. Set FLASHINFER_ROCM_ARCH_LIST to the architecture you " + "are building for.", + fallback, + ) + return fallback + + def get_rocm_home(): """ Get the ROCM_HOME directory from environment variables or default path. @@ -190,7 +251,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: Args: arch_list: Comma-separated list of architectures (e.g., "gfx942,gfx90a"). - If None, reads from FLASHINFER_ROCM_ARCH_LIST env var or defaults to "gfx942" + If None, resolved by :func:`resolve_target_archs`. verbose: Whether to print validation messages Returns: @@ -199,7 +260,6 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: Raises: RuntimeError: If ROCm not found or architectures not supported """ - import os # ROCm compatibility matrix: version -> supported gfx architectures # Refer: https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html @@ -234,7 +294,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: # Get architecture list from parameter, env var, or default if arch_list is None: - arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942") + arch_list = resolve_target_archs() # Validate system has ROCm installed system_rocm_version = get_system_rocm_version() @@ -315,11 +375,10 @@ def validate_flashinfer_rocm_arch( Raises: RuntimeError: If any validation step fails with clear error message """ - import os # Get architecture list from parameter, env var, or default if arch_list is None: - arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942") + arch_list = resolve_target_archs() # Step 1: Validate against system ROCm version (reuse existing logic) validated_arch_list = validate_rocm_arch(arch_list=arch_list, verbose=verbose) diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 75ea294247..00c159e659 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -23,6 +23,7 @@ get_supported_device_indices, get_system_rocm_version_from_hipconfig, is_therock_build, + resolve_target_archs, validate_flashinfer_rocm_arch, validate_rocm_arch, ) @@ -137,7 +138,87 @@ def test_returns_none_on_timeout(self): assert get_system_rocm_version_from_hipconfig() is None -# validate_rocm_arch +# resolve_target_archs +class TestResolveTargetArchs: + """The single resolver every build path now consults. + + Before it existed, three call sites answered "what are we building for" + independently and could disagree: on a gfx950 host, + ``validate_flashinfer_rocm_arch(arch_list=None)`` returned ``{"gfx942"}`` + while ``CompilationContext`` emitted ``--offload-arch=gfx950``. + """ + + def _agents(self, *archs): + return patch( + "flashinfer.hip_utils.rocminfo_gpu_agents", + return_value=tuple((arch, "") for arch in archs), + ) + + def test_explicit_argument_wins(self, monkeypatch): + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx950") + with self._agents("gfx950"): + assert resolve_target_archs("gfx942") == "gfx942" + + def test_env_var_beats_detection(self, monkeypatch): + """An explicit request is honoured even when it is not what is plugged in + -- cross-compiling for the other architecture must stay possible.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942") + with self._agents("gfx950"): + assert resolve_target_archs() == "gfx942" + + def test_detects_the_running_device(self, monkeypatch): + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents("gfx950"): + assert resolve_target_archs() == "gfx950" + + def test_detects_every_distinct_supported_arch(self, monkeypatch): + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents("gfx942", "gfx950", "gfx942"): + assert resolve_target_archs() == "gfx942,gfx950" + + def test_ignores_unsupported_agents(self, monkeypatch): + """An integrated GPU alongside the dGPU must not widen the target set.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents("gfx1035", "gfx950"): + assert resolve_target_archs() == "gfx950" + + def test_no_device_builds_for_everything_supported(self, monkeypatch, caplog): + """A GPU-less build host gets a fat build, not a guess. + + Picking one architecture here is what produced a silently gfx942-only + artifact on a machine that could not confirm anything. Building for all + supported architectures is slower but correct wherever it lands, and the + warning tells the operator how to make it cheaper. + """ + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents(), caplog.at_level("WARNING"): + assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + assert "FLASHINFER_ROCM_ARCH_LIST" in caplog.text + + def test_agrees_with_the_compilation_context(self, monkeypatch): + """The property the whole change exists for: the validator and the thing + that emits --offload-arch must not disagree.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + from flashinfer.compilation_context_hip import CompilationContext + + with ( + self._agents("gfx950"), + patch("flashinfer.hip_utils.get_system_rocm_version", return_value="7.1.0"), + ): + _, validated = validate_flashinfer_rocm_arch( + arch_list=None, torch_cpp_ext_module=_FakeCppExt(), verbose=False + ) + assert validated == CompilationContext().TARGET_ROCM_ARCHS + + +class _FakeCppExt: + """Stands in for torch.utils.cpp_extension: claims both archs are built in.""" + + @staticmethod + def _get_rocm_arch_flags(): + return [f"--offload-arch={a}" for a in FLASHINFER_SUPPORTED_ROCM_ARCHS] + + class TestValidateRocmArch: def _patch_rocm_version(self, version): return patch( @@ -188,10 +269,24 @@ def test_reads_arch_from_env_when_none_given(self, monkeypatch): with self._patch_rocm_version("7.1.0"): assert validate_rocm_arch(arch_list=None) == "gfx942" - def test_defaults_to_gfx942_when_no_env_and_no_arg(self, monkeypatch): + def test_falls_back_to_the_running_device_not_a_hard_coded_arch(self, monkeypatch): + """With no argument and no env var, follow the hardware. + + This used to assert ``== "gfx942"``, encoding the literal that made + ``validate_flashinfer_rocm_arch(arch_list=None)`` answer ``gfx942`` on a + gfx950 device while CompilationContext compiled for gfx950. A test that + pins a wrong constant is how the constant survives, so it is now pinned + to the detected architecture instead. + """ monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) - with self._patch_rocm_version("7.1.0"): - assert validate_rocm_arch(arch_list=None) == "gfx942" + with ( + self._patch_rocm_version("7.1.0"), + patch( + "flashinfer.hip_utils.rocminfo_gpu_agents", + return_value=(("gfx950", "AMD Instinct MI350X"),), + ), + ): + assert validate_rocm_arch(arch_list=None) == "gfx950" def test_verbose_prints_message(self, capsys): with self._patch_rocm_version("7.1.0"): From 1937e0c5ecc1ed75140ccef9ae37daf55ee6dfc6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 12:10:24 -0400 Subject: [PATCH 02/16] fix(rocm): canonicalize the resolved arch list, and make the agreement test hermetic Addresses both review comments on #289. - resolve_target_archs() returned the caller/env string verbatim. Now that it is the single source of truth that is a hard failure, not untidiness: the validators split on "," only and match tokens against FLASHINFER_SUPPORTED_ROCM_ARCHS verbatim, so FLASHINFER_ROCM_ARCH_LIST=gfx950:sramecc+ -> ['gfx950:sramecc+'] unsupported FLASHINFER_ROCM_ARCH_LIST=gfx942;gfx950 -> ['gfx942;gfx950'] unsupported FLASHINFER_ROCM_ARCH_LIST=gfx942,,gfx942 -> ['gfx942','','gfx942'] unsupported '' and validate_flashinfer_rocm_arch raises "does not support any of the requested ROCm architectures". ';' matters specifically because jit/aiter_source.py already documents it for this same variable, and aot_hip.py writes this resolver's output back into that env var -- so the two consumers were disagreeing about their own input format. _canonical_arch_list normalizes syntax only: accepts ',' or ';', strips qualifiers via normalize_arch, drops empties, dedupes preserving first-seen order. Unknown architectures pass through so the validators can still report them; dropping one here would turn a clear error into a build that quietly targets less than was asked for. A value that normalizes away entirely (";;") falls through to detection rather than returning "". - test_agrees_with_the_compilation_context compared a _FakeCppExt-fed validator against a CompilationContext that validates against the *real* torch, so the assertion depended on the installed wheel. Both sides now see the same view. Verified rather than assumed, by simulating an arch-specific wheel (_get_rocm_arch_flags -> gfx942 only) via a pytest plugin: before: RuntimeError: PyTorch does not support the following architectures: --offload-arch=gfx950 -> FAILED after: passes The wheel on this box is a fat build advertising gfx950, which is why the fragility was latent here. 11 tests added for the canonicalization. gfx942 behaviour is unchanged. --- flashinfer/hip_utils.py | 40 +++++++++++++++++- tests/rocm_tests/test_hip_utils.py | 66 ++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 2 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 9016dcb49a..fd885ca2fb 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -16,6 +16,34 @@ FLASHINFER_SUPPORTED_ROCM_ARCHS = ["gfx942", "gfx950"] +def _canonical_arch_list(raw: str) -> str: + """``"gfx950:sramecc+; gfx942,,gfx942"`` -> ``"gfx950,gfx942"``. + + Caller- and environment-supplied lists arrive in whatever shape the operator + typed. The validators below split on ``","`` only and compare tokens against + :data:`FLASHINFER_SUPPORTED_ROCM_ARCHS` verbatim, so an unnormalized value is + not merely untidy -- ``"gfx942;gfx950"`` becomes the single token + ``"gfx942;gfx950"``, matches nothing, and + :func:`validate_flashinfer_rocm_arch` raises "FlashInfer does not support any + of the requested ROCm architectures". ``";"`` is worth accepting because + ``jit/aiter_source.py`` already documents it for this same variable, and the + two must not disagree about their own env var. + + Only *syntax* is normalized. Unknown architectures are passed through so the + validators can report them; silently dropping one here would turn a clear + error into a build that quietly targets less than was asked for. + + Order is preserved (first occurrence wins) rather than sorted: it is the + caller's stated preference, and it is what ends up on the hipcc command line. + """ + seen = [] + for token in raw.replace(";", ",").split(","): + arch = normalize_arch(token) + if arch and arch not in seen: + seen.append(arch) + return ",".join(seen) + + def resolve_target_archs(arch_list: str = None) -> str: """Return the architectures to build for, as a comma-separated string. @@ -46,12 +74,20 @@ def resolve_target_archs(arch_list: str = None) -> str: """ import os + # Canonicalize the two operator-supplied paths. A value that normalizes away + # entirely (";;", whitespace) falls through to detection rather than + # returning "", which would otherwise reach the validators as a single empty + # token and fail as an unsupported architecture. if arch_list: - return arch_list + canonical = _canonical_arch_list(arch_list) + if canonical: + return canonical from_env = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") if from_env: - return from_env + canonical = _canonical_arch_list(from_env) + if canonical: + return canonical detected = sorted( { diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 00c159e659..18959b43c2 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -166,6 +166,57 @@ def test_env_var_beats_detection(self, monkeypatch): with self._agents("gfx950"): assert resolve_target_archs() == "gfx942" + @pytest.mark.parametrize( + "raw,expected", + [ + # Qualifiers: what torch's gcnArchName looks like, so an operator + # copying from `rocminfo` or a torch error message pastes this shape. + ("gfx950:sramecc+:xnack-", "gfx950"), + # ';' is documented for this same variable by jit/aiter_source.py. + # The two consumers must not disagree about their own env var. + ("gfx942;gfx950", "gfx942,gfx950"), + ("gfx942; gfx950", "gfx942,gfx950"), + # Empty tokens would otherwise reach the validators as "" and be + # reported as an unsupported architecture. + ("gfx942,,gfx950", "gfx942,gfx950"), + ("gfx942, gfx950 ", "gfx942,gfx950"), + # Duplicates collapse; first occurrence sets the order, which is what + # lands on the hipcc command line. + ("gfx950,gfx942,gfx950", "gfx950,gfx942"), + # Already canonical input must round-trip untouched. + ("gfx942,gfx950", "gfx942,gfx950"), + ], + ) + def test_env_var_is_canonicalized(self, monkeypatch, raw, expected): + """The validators split on ',' only and match tokens verbatim, so an + unnormalized value is a hard failure, not an untidiness: "gfx942;gfx950" + arrives as one token, matches nothing, and validate_flashinfer_rocm_arch + raises "does not support any of the requested ROCm architectures".""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", raw) + with self._agents("gfx950"): + assert resolve_target_archs() == expected + + def test_explicit_argument_is_canonicalized(self, monkeypatch): + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents("gfx950"): + assert resolve_target_archs("gfx942:xnack-;gfx950") == "gfx942,gfx950" + + def test_unknown_archs_survive_normalization(self, monkeypatch): + """Only syntax is normalized. Dropping an unrecognized arch here would + turn the validators' clear error into a build that quietly targets less + than was asked for.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx900:xnack-;gfx942") + with self._agents("gfx950"): + assert resolve_target_archs() == "gfx900,gfx942" + + @pytest.mark.parametrize("raw", [";;", " ", ",", " ; , "]) + def test_a_value_that_normalizes_away_falls_through(self, monkeypatch, raw): + """Returning "" would reach the validators as a single empty token and + fail as an unsupported architecture; detection is the better answer.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", raw) + with self._agents("gfx950"): + assert resolve_target_archs() == "gfx950" + def test_detects_the_running_device(self, monkeypatch): monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) with self._agents("gfx950"): @@ -199,11 +250,26 @@ def test_agrees_with_the_compilation_context(self, monkeypatch): """The property the whole change exists for: the validator and the thing that emits --offload-arch must not disagree.""" monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + import torch.utils.cpp_extension as torch_cpp_ext + from flashinfer.compilation_context_hip import CompilationContext + # Both sides must see the same PyTorch. CompilationContext imports the + # real torch.utils.cpp_extension and validates against it, while the + # direct call below is handed _FakeCppExt -- so without this patch the + # assertion depends on the installed wheel. The wheel here is a fat build + # advertising gfx950, which is why that went unnoticed; an arch-specific + # build (gfx942-only) would make CompilationContext() raise "PyTorch does + # not support the following architectures" and fail this test for a + # reason that has nothing to do with resolver agreement. with ( self._agents("gfx950"), patch("flashinfer.hip_utils.get_system_rocm_version", return_value="7.1.0"), + patch.object( + torch_cpp_ext, + "_get_rocm_arch_flags", + _FakeCppExt._get_rocm_arch_flags, + ), ): _, validated = validate_flashinfer_rocm_arch( arch_list=None, torch_cpp_ext_module=_FakeCppExt(), verbose=False From 4d1335d327c0a45e974aed22b6b5f45e370e9cb7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 12:24:04 -0400 Subject: [PATCH 03/16] fix(rocm): validate the arch list before publishing it to the environment Addresses the suppressed comment in Copilot review 4984984345 on #289. compile_and_package_modules set FLASHINFER_ROCM_ARCH_LIST and then validated it, so a CompilationContext() raise left the variable set for whatever ran next in the same process. Publishing it is deliberate -- the AITER shim reads the build architecture from there (jit/aiter_source.py) and an AOT build has no other channel -- but it is a side effect that outlives the call, so it should only happen once the list is known good. Reordering is behaviour-preserving on success: CompilationContext() re-resolves through resolve_target_archs() and nothing can change between the two calls, so it validates exactly the list published afterwards. On the failure path the environment is now left as it was found. One correction to the review's wording: the leaked value is not "invalid" -- resolve_target_archs() returns a canonical list. It is a *valid* list that failed validation against this ROCm version or PyTorch build, which is a different thing and is why the leak is subtle rather than obvious. The regression test was wrong on its first attempt and is worth flagging: it seeded FLASHINFER_ROCM_ARCH_LIST with the value the resolver would return, so the buggy write was a no-op and the test passed with the bug present. It now starts from the variable unset with detection patched, and was verified to fail without the reorder: AssertionError: assert 'FLASHINFER_ROCM_ARCH_LIST' not in environ({... 'FLASHINFER_ROCM_ARCH_LIST': 'gfx950'}) --- flashinfer/aot_hip.py | 7 ++++- tests/rocm_tests/test_aot_hip.py | 48 ++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 5c4ee8f18c..5914321394 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -224,9 +224,14 @@ def compile_and_package_modules( from .compilation_context_hip import CompilationContext from .hip_utils import resolve_target_archs + # Validate before publishing, so a failed build leaves the environment as it + # found it. CompilationContext() re-resolves through the same function and + # the inputs cannot change in between, so it validates exactly the list + # published below -- the order costs nothing and stops a raise here from + # leaving FLASHINFER_ROCM_ARCH_LIST set for whatever runs next in-process. rocm_arch_list = resolve_target_archs() - os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list CompilationContext() # validates the resolved list, raising on a bad one + os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list if verbose: print(f"Target ROCm architectures: {rocm_arch_list}") diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index c3873f2078..30c88a1071 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -11,6 +11,7 @@ 3. .so files are created and can be loaded """ +import os import shutil import tempfile from pathlib import Path @@ -139,6 +140,53 @@ def test_compile_and_package_minimal(): shutil.rmtree(out_dir, ignore_errors=True) +def test_failed_validation_leaves_the_environment_alone(monkeypatch): + """A raise mid-build must not leave FLASHINFER_ROCM_ARCH_LIST behind. + + compile_and_package_modules publishes the resolved list into the process + environment on purpose -- the AITER shim reads it from there + (jit/aiter_source.py) and an AOT build has no other channel to reach it. + That side effect outlives the call, so it must only happen once the list is + known good; otherwise a build that dies on an unsupported ROCm version + silently repoints whatever runs next in the same process. + + The variable starts *unset* and resolution comes from detection, so the + published value differs from the starting state. Seeding it with the value + the resolver would return makes the write a no-op and the test vacuous -- + it then passes with the bug present. + """ + import flashinfer.aot_hip as aot_hip + + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setattr( + "flashinfer.hip_utils.rocminfo_gpu_agents", + lambda: (("gfx950", ""),), + ) + + class _Boom: + def __init__(self): + raise RuntimeError("ROCm version 0.0 is not recognized") + + monkeypatch.setattr("flashinfer.compilation_context_hip.CompilationContext", _Boom) + + with pytest.raises(RuntimeError, match="not recognized"): + aot_hip.compile_and_package_modules( + out_dir=None, + build_dir=Path(tempfile.mkdtemp()), + project_root=Path(__file__).parent.parent, + config={ + "fa2_head_dim": [(128, 128)], + "f16_dtype": [torch.float16], + "use_sliding_window": [False], + "use_logits_soft_cap": [False], + }, + verbose=False, + skip_prebuilt=True, + ) + + assert "FLASHINFER_ROCM_ARCH_LIST" not in os.environ + + def test_module_naming_convention(): """Test that generated module names follow expected conventions.""" f16_dtype = [torch.float16] From 6e8773c5590b7a938f5bdff758f105e798f34f94 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 12:39:25 -0400 Subject: [PATCH 04/16] fix(rocm): publish the arch list before validating, and restore it on failure Addresses the suppressed comment in Copilot review 4985096141 on #289, which caught a regression introduced by the previous commit. Validating before publishing made CompilationContext() re-resolve from an unset environment. rocminfo_gpu_agents() is not cached -- only get_supported_device_indices and get_physical_card_device_indices are -- so that is a second detection pass, and on a GPU-less host a second "no supported AMD GPU detected" warning. Measured: two resolves emit two warnings. Publishing first and restoring in an except clause gets both properties at once: validation sees exactly the list that was resolved (one detection, one warning), and a raise leaves the variable as it was found, including the case where it was previously unset. The test now pins both halves, since the leak fix alone passed without the single-resolution property: _Boom records what os.environ held when validation ran, so a return to validate-then-publish fails on `_Boom.seen == "gfx950"` rather than quietly reintroducing the duplicate detection. --- flashinfer/aot_hip.py | 23 +++++++++++++++++------ tests/rocm_tests/test_aot_hip.py | 10 ++++++++++ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 5914321394..477d98e259 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -224,14 +224,25 @@ def compile_and_package_modules( from .compilation_context_hip import CompilationContext from .hip_utils import resolve_target_archs - # Validate before publishing, so a failed build leaves the environment as it - # found it. CompilationContext() re-resolves through the same function and - # the inputs cannot change in between, so it validates exactly the list - # published below -- the order costs nothing and stops a raise here from - # leaving FLASHINFER_ROCM_ARCH_LIST set for whatever runs next in-process. + # Publish first so CompilationContext() -- which resolves through + # resolve_target_archs() itself -- validates exactly this list rather than + # repeating the work. rocminfo_gpu_agents() is not cached, so simply + # validating before publishing would re-run detection and, on a GPU-less + # host, emit the "no supported AMD GPU detected" warning a second time. + # + # Restore on failure so a build that dies here does not leave the variable + # pointing somewhere for whatever runs next in-process. rocm_arch_list = resolve_target_archs() - CompilationContext() # validates the resolved list, raising on a bad one + previous_arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list + try: + CompilationContext() # validates the resolved list, raising on a bad one + except BaseException: + if previous_arch_list is None: + os.environ.pop("FLASHINFER_ROCM_ARCH_LIST", None) + else: + os.environ["FLASHINFER_ROCM_ARCH_LIST"] = previous_arch_list + raise if verbose: print(f"Target ROCm architectures: {rocm_arch_list}") diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index 30c88a1071..5d1501fd8c 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -164,7 +164,12 @@ def test_failed_validation_leaves_the_environment_alone(monkeypatch): ) class _Boom: + """Records what validation actually saw, then fails the build.""" + + seen = "unset" + def __init__(self): + _Boom.seen = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") raise RuntimeError("ROCm version 0.0 is not recognized") monkeypatch.setattr("flashinfer.compilation_context_hip.CompilationContext", _Boom) @@ -184,6 +189,11 @@ def __init__(self): skip_prebuilt=True, ) + # Validation must see the resolved list, not an unset variable. Otherwise + # CompilationContext re-resolves from scratch -- rocminfo_gpu_agents() is + # not cached, so that is a second detection pass and, on a GPU-less host, a + # second "no supported AMD GPU detected" warning. + assert _Boom.seen == "gfx950" assert "FLASHINFER_ROCM_ARCH_LIST" not in os.environ From 7bdc2c9f7db400bad238504463204691cfa5b5aa Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 12:54:26 -0400 Subject: [PATCH 05/16] fix(rocm): publish the validated arch list, not the resolved superset Addresses both suppressed comments in Copilot review 4985222187 on #289. - Validation filters rather than raises. An architecture FlashInfer cannot serve is dropped with a warnings.warn and the build continues, provided at least one survives. Measured: validate_flashinfer_rocm_arch("gfx900,gfx942") -> no exception, arch_flags ['--offload-arch=gfx942'], set {'gfx942'} so FLASHINFER_ROCM_ARCH_LIST could advertise gfx900 while the packaged kernels were compiled only for gfx942. The AITER shim resolves its own build target from that variable (jit/aiter_source.py), so it would build for an architecture nothing else in the package targets -- the precise divergence this PR exists to remove, reintroduced one layer up. The resolved list is still published before validation, so CompilationContext does not re-run detection; the validated list replaces it afterwards. Taken from arch_flags rather than TARGET_ROCM_ARCHS because the latter is a set and order is meaningful, both on the hipcc command line and to AITER. - The new failure-path test leaked a tempfile.mkdtemp() directory. Switched to the tmp_path fixture, which pytest cleans up, rather than adding a finally block. Verified the guard is real, not vacuous -- without the republish: - gfx950 + gfx900,gfx950 --- flashinfer/aot_hip.py | 17 +++++++++++++- tests/rocm_tests/test_aot_hip.py | 38 ++++++++++++++++++++++++++++++-- 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 477d98e259..742d3d5e9a 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -236,13 +236,28 @@ def compile_and_package_modules( previous_arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list try: - CompilationContext() # validates the resolved list, raising on a bad one + context = CompilationContext() # raises if nothing in the list is usable except BaseException: if previous_arch_list is None: os.environ.pop("FLASHINFER_ROCM_ARCH_LIST", None) else: os.environ["FLASHINFER_ROCM_ARCH_LIST"] = previous_arch_list raise + + # Republish what will *actually* be compiled. Validation drops architectures + # it cannot serve with a warning rather than raising, provided at least one + # survives -- "gfx900,gfx942" validates down to "gfx942" -- so the resolved + # list can be a strict superset of the built one. Leaving the superset in + # the environment is the exact failure this PR exists to remove: the AITER + # shim resolves its own target from this variable, so it would build for an + # architecture the packaged kernels were not compiled for. + # + # Taken from arch_flags rather than TARGET_ROCM_ARCHS because that is a set; + # order is meaningful here, both on the hipcc command line and to AITER. + rocm_arch_list = ",".join( + flag.removeprefix("--offload-arch=") for flag in context.arch_flags + ) + os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list if verbose: print(f"Target ROCm architectures: {rocm_arch_list}") diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index 5d1501fd8c..7135bb3423 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -14,6 +14,7 @@ import os import shutil import tempfile +import warnings from pathlib import Path import pytest @@ -140,7 +141,7 @@ def test_compile_and_package_minimal(): shutil.rmtree(out_dir, ignore_errors=True) -def test_failed_validation_leaves_the_environment_alone(monkeypatch): +def test_failed_validation_leaves_the_environment_alone(monkeypatch, tmp_path): """A raise mid-build must not leave FLASHINFER_ROCM_ARCH_LIST behind. compile_and_package_modules publishes the resolved list into the process @@ -177,7 +178,7 @@ def __init__(self): with pytest.raises(RuntimeError, match="not recognized"): aot_hip.compile_and_package_modules( out_dir=None, - build_dir=Path(tempfile.mkdtemp()), + build_dir=tmp_path, project_root=Path(__file__).parent.parent, config={ "fa2_head_dim": [(128, 128)], @@ -197,6 +198,39 @@ def __init__(self): assert "FLASHINFER_ROCM_ARCH_LIST" not in os.environ +def test_environment_gets_the_validated_list_not_the_resolved_one( + monkeypatch, tmp_path +): + """Validation filters silently, so publishing the resolved list is not enough. + + An architecture FlashInfer cannot serve is dropped with a warning rather + than an exception, provided at least one survives -- "gfx900,gfx950" + validates down to "gfx950". The AITER shim resolves its own build target + from FLASHINFER_ROCM_ARCH_LIST (jit/aiter_source.py), so leaving the wider + list there would have it build for an architecture the packaged kernels + were never compiled for. That divergence is the whole point of this PR. + """ + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx900,gfx950") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + compile_and_package_modules( + out_dir=None, + build_dir=tmp_path, + project_root=Path(__file__).parent.parent, + config={ + "fa2_head_dim": [(128, 128)], + "f16_dtype": [torch.float16], + "use_sliding_window": [False], + "use_logits_soft_cap": [False], + }, + verbose=False, + skip_prebuilt=True, + ) + + assert os.environ["FLASHINFER_ROCM_ARCH_LIST"] == "gfx950" + + def test_module_naming_convention(): """Test that generated module names follow expected conventions.""" f16_dtype = [torch.float16] From b2f586d90a57952f4ee39250bde7c93d56f2bc12 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 13:52:13 -0400 Subject: [PATCH 06/16] perf(rocm): cache the rocminfo probe; sharpen the fallback warning and a stale test Addresses three of the six findings in Copilot review 4985367370 on #289. The other three are design-level and tracked separately. - rocminfo_gpu_agents() is now cached. gen_jit_spec() calls check_rocm_arch() for every module it builds, and that reaches the probe through resolve_target_archs(), so a multi-op process was paying one rocminfo subprocess -- with a 10 s timeout -- per operation, and a GPU-less host repeated the fallback warning each time. The _auto_detect_archs() this PR removed avoided that by going through the already-cached get_supported_device_indices(), so this is a regression the PR introduced. Caching the probe itself also makes that function's "rocminfo is invoked at most once per process" docstring true globally rather than only on its path. TestGetSupportedDeviceIndices already cleared the derived cache per test; it now clears this one too. Without that the first test in the class pinned the probe and the remaining five asserted against it instead of their own patched subprocess output -- which is how they failed when the cache was added. - The GPU-less fallback warning said FLASHINFER_ROCM_ARCH_LIST "is unset" even when it was set to something that normalized away (";;", whitespace), which is reachable through the fall-through added earlier in this branch. It now distinguishes unset from "names no architecture", so the operator is not sent looking for a variable that is already there. - test_defaults_to_gfx942_when_no_env_no_arg stubbed validate_rocm_arch to return "gfx942" whatever it was handed, so the resolved list never reached the assertion: the hard-coded default this PR exists to remove could have come back with the test still green. The stub now echoes its argument and the detected architecture is what is checked. --- flashinfer/hip_utils.py | 28 ++++++++++++++++++++---- tests/rocm_tests/test_hip_utils.py | 34 +++++++++++++++++++++++++----- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index fd885ca2fb..1a66dfdb46 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -100,11 +100,19 @@ def resolve_target_archs(arch_list: str = None) -> str: return ",".join(detected) fallback = ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + # Distinguish the two ways of getting here. Telling an operator who *did* + # set the variable to set it sends them looking for a variable that is + # already there; the actual problem is that its value held no architecture. + how = ( + "FLASHINFER_ROCM_ARCH_LIST is unset" + if from_env is None + else f"FLASHINFER_ROCM_ARCH_LIST={from_env!r} names no architecture" + ) logger.warning( - "No supported AMD GPU detected and FLASHINFER_ROCM_ARCH_LIST is unset; " - "building for every supported architecture (%s). This is slower than " - "targeting one. Set FLASHINFER_ROCM_ARCH_LIST to the architecture you " - "are building for.", + "No supported AMD GPU detected and %s; building for every supported " + "architecture (%s). This is slower than targeting one. Set " + "FLASHINFER_ROCM_ARCH_LIST to the architecture you are building for.", + how, fallback, ) return fallback @@ -485,10 +493,22 @@ def get_available_gpu_count() -> int: return torch.cuda.device_count() +@functools.cache def rocminfo_gpu_agents() -> tuple: """ Return ``(arch, marketing_name)`` for each GPU agent rocminfo reports. + Cached for the process. ``gen_jit_spec`` calls ``check_rocm_arch()`` for + every operation it builds, and that reaches here through + ``resolve_target_archs()``, so without this a multi-op process pays one + ``rocminfo`` subprocess -- and its 10-second timeout -- per module, and a + GPU-less host repeats the fallback warning each time. The prior + ``_auto_detect_archs()`` avoided this by going through the already-cached + ``get_supported_device_indices()``; caching the probe itself restores that + and makes the "rocminfo is invoked at most once per process" claim in + ``get_supported_device_indices`` true globally rather than only on its own + path. The hardware cannot change under a running process. + Uses rocminfo (subprocess) rather than torch.cuda so this is safe to call before the HIP runtime is initialized (e.g. before HIP_VISIBLE_DEVICES is set in xdist workers). rocminfo enumerates GPU agents in the same order as the diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 18959b43c2..98b7677112 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -24,6 +24,7 @@ get_system_rocm_version_from_hipconfig, is_therock_build, resolve_target_archs, + rocminfo_gpu_agents, validate_flashinfer_rocm_arch, validate_rocm_arch, ) @@ -445,11 +446,25 @@ def test_reads_arch_from_env_when_none_given(self, monkeypatch): flags, arch_set = validate_flashinfer_rocm_arch(arch_list=None) assert arch_set == {"gfx942"} - def test_defaults_to_gfx942_when_no_env_no_arg(self, monkeypatch): + def test_no_env_no_arg_follows_the_detected_device(self, monkeypatch): + """The wrapper must pass the *resolved* list through, not a constant. + + This previously stubbed validate_rocm_arch to return "gfx942" whatever + it was given and asserted "gfx942" back, so the resolution never reached + the assertion -- the same hard-coded default this PR removes could have + come back and the test would still have passed. The stub now echoes its + argument, so the detected architecture is what is being checked. + """ monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) - with self._patch_validate_rocm_arch("gfx942"): - flags, arch_set = validate_flashinfer_rocm_arch(arch_list=None) - assert arch_set == {"gfx942"} + monkeypatch.setattr( + "flashinfer.hip_utils.rocminfo_gpu_agents", lambda: (("gfx950", ""),) + ) + with patch( + "flashinfer.hip_utils.validate_rocm_arch", + side_effect=lambda arch_list, verbose=False: arch_list, + ): + _, arch_set = validate_flashinfer_rocm_arch(arch_list=None) + assert arch_set == {"gfx950"} def test_verbose_prints_message(self, capsys): with self._patch_validate_rocm_arch("gfx942"): @@ -519,9 +534,18 @@ class TestGetSupportedDeviceIndices: """Each test clears the functools.cache to avoid cross-test contamination.""" def setup_method(self): - get_supported_device_indices.cache_clear() + self._clear_caches() def teardown_method(self): + self._clear_caches() + + @staticmethod + def _clear_caches(): + # rocminfo_gpu_agents is cached too, so the derived cache alone is not + # enough: the first test in the class would otherwise pin the probe + # result and the rest would assert against it instead of their own + # patched subprocess output. + rocminfo_gpu_agents.cache_clear() get_supported_device_indices.cache_clear() def _run_result(self, stdout, returncode=0): From b3fd02395174d3a49c6c58c2d5cde1fc170e95b2 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 14:07:53 -0400 Subject: [PATCH 07/16] fix(rocm): repair the comments the cache invalidated, and widen probe failure handling Addresses seven of the eight suppressed comments in Copilot review 4985855053 on #289. Four of them are defects the previous commit introduced. Self-inflicted by adding @functools.cache to rocminfo_gpu_agents: - The docstring still carried "Not cached: the caller decides" three paragraphs below the new "Cached for the process". Removed. - aot_hip's rationale for publishing before validating said it avoided "a second detection pass". True when written, false once the probe was cached. Rewritten to the reason that survives: the two agree by construction rather than by coincidence. - The same stale reasoning in the test comment, likewise rewritten. Also self-inflicted, and the same bug flagged earlier on a different test: - test_environment_gets_the_validated_list_not_the_resolved_one hard-coded gfx950 and validated against the real PyTorch, so on an arch-specific gfx942 wheel it would fail for reasons unrelated to the republish. Stubs _get_rocm_arch_flags now. Verified with a plugin simulating a gfx942-only wheel: this test, test_failed_validation_leaves_the_environment_alone and test_agrees_with_the_compilation_context all pass under it. Pre-existing, surfaced by this PR changing the behaviour they describe: - jit/core.py's "defaults to gfx942" comment and amd-flashinfer-jit-cache's README both still documented the hard-coded default this PR removes. - rocminfo_gpu_agents caught only FileNotFoundError and TimeoutExpired, so a present-but-unexecutable rocminfo (PermissionError) aborted resolution instead of reaching the documented GPU-less fallback. The removed _auto_detect_archs() swallowed these; catching OSError restores that. The eighth -- the module-global JIT compilation context being a different object from the one validated here -- is deferred with the other two design-level findings; see the thread replies. --- amd-flashinfer-jit-cache/README.md | 6 +++--- flashinfer/aot_hip.py | 6 +++--- flashinfer/hip_utils.py | 12 ++++++++---- flashinfer/jit/core.py | 2 +- tests/rocm_tests/test_aot_hip.py | 20 ++++++++++++++++---- 5 files changed, 31 insertions(+), 15 deletions(-) diff --git a/amd-flashinfer-jit-cache/README.md b/amd-flashinfer-jit-cache/README.md index c2bb2f47a9..a895fb0c64 100644 --- a/amd-flashinfer-jit-cache/README.md +++ b/amd-flashinfer-jit-cache/README.md @@ -16,7 +16,7 @@ pip install amd-flashinfer amd-flashinfer-jit-cache ## Architecture Support -This package is built specifically for the **AMD MI300 series (gfx942)** architecture. +This package is built for whichever architectures the build host resolves to. `flashinfer.aot_hip` asks `hip_utils.resolve_target_archs()`, which takes `FLASHINFER_ROCM_ARCH_LIST` if set, else the supported GPUs actually present, else every supported architecture (gfx942 and gfx950) with a warning. Set the variable to pin a single target. Check the package version and tags to ensure compatibility with your GPU architecture. @@ -32,12 +32,12 @@ python -m build --wheel The build process will: 1. Generate kernel specifications using `flashinfer.aot_hip` -2. Compile kernels for the gfx942 architecture +2. Compile kernels for the resolved architectures (see Architecture Support) 3. Package compiled `.so` files into the wheel ## Environment Variables -- `FLASHINFER_ROCM_ARCH_LIST`: Target architecture (default: "gfx942") +- `FLASHINFER_ROCM_ARCH_LIST`: Target architectures, comma- or semicolon-separated. Unset means detect, then fall back to all supported. - `HIP_PATH`: Path to ROCm/HIP installation (auto-detected if not set) ## License diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 742d3d5e9a..931f271de0 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -226,9 +226,9 @@ def compile_and_package_modules( # Publish first so CompilationContext() -- which resolves through # resolve_target_archs() itself -- validates exactly this list rather than - # repeating the work. rocminfo_gpu_agents() is not cached, so simply - # validating before publishing would re-run detection and, on a GPU-less - # host, emit the "no supported AMD GPU detected" warning a second time. + # re-deriving one. rocminfo_gpu_agents() is cached now, so this is no longer + # about avoiding a second subprocess; it is about the two agreeing by + # construction instead of by coincidence, which is the whole point of the PR. # # Restore on failure so a build that dies here does not leave the variable # pointing somewhere for whatever runs next in-process. diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 1a66dfdb46..914601b478 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -518,9 +518,6 @@ def rocminfo_gpu_agents() -> tuple: "Marketing Name" (e.g. "AMD Instinct MI350X") or "" when absent. CPU agents also carry a Marketing Name, so agents are filtered on ``Device Type: GPU``. - Not cached: the caller decides. ``get_supported_device_indices`` caches its - own result, and one-shot consumers pay a single subprocess. - Returns: tuple[tuple[str, str], ...]: One (arch, marketing_name) pair per GPU agent. Empty if rocminfo is unavailable. @@ -534,7 +531,14 @@ def rocminfo_gpu_agents() -> tuple: ) if result.returncode != 0: return () - except (FileNotFoundError, subprocess.TimeoutExpired): + except (OSError, subprocess.TimeoutExpired): + # OSError rather than FileNotFoundError alone: a rocminfo that exists + # but cannot be executed raises PermissionError, and a container with + # no /dev/kfd can raise other OSErrors. Every one of them means the + # same thing here -- we could not ask the hardware -- and the caller's + # documented answer to that is the GPU-less fallback, not a traceback + # out of resolve_target_archs(). The removed _auto_detect_archs() + # swallowed these; this keeps that behaviour. return () agents = [] diff --git a/flashinfer/jit/core.py b/flashinfer/jit/core.py index af94cf6026..ac4a432898 100644 --- a/flashinfer/jit/core.py +++ b/flashinfer/jit/core.py @@ -145,7 +145,7 @@ def check_rocm_arch(): try: validate_flashinfer_rocm_arch( - arch_list=None, # Uses FLASHINFER_ROCM_ARCH_LIST env or defaults to gfx942 + arch_list=None, # resolve_target_archs(): env, else detected, else all supported torch_cpp_ext_module=torch_cpp_ext, verbose=False, ) diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index 7135bb3423..24babb3ca3 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -190,10 +190,9 @@ def __init__(self): skip_prebuilt=True, ) - # Validation must see the resolved list, not an unset variable. Otherwise - # CompilationContext re-resolves from scratch -- rocminfo_gpu_agents() is - # not cached, so that is a second detection pass and, on a GPU-less host, a - # second "no supported AMD GPU detected" warning. + # Validation must see the resolved list, not an unset variable, so that the + # list validated and the list published are the same object rather than two + # independent derivations that happen to agree. assert _Boom.seen == "gfx950" assert "FLASHINFER_ROCM_ARCH_LIST" not in os.environ @@ -210,7 +209,20 @@ def test_environment_gets_the_validated_list_not_the_resolved_one( list there would have it build for an architecture the packaged kernels were never compiled for. That divergence is the whole point of this PR. """ + import torch.utils.cpp_extension as torch_cpp_ext + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx900,gfx950") + # Stub PyTorch's advertised architectures. CompilationContext validates + # against the real wheel, so without this the test asserts something about + # the installed torch rather than about the republish: on an arch-specific + # gfx942 build, gfx900 is filtered and then step 3 raises because gfx950 is + # not advertised. The wheel here is a fat build, which is exactly why that + # would have gone unnoticed. + monkeypatch.setattr( + torch_cpp_ext, + "_get_rocm_arch_flags", + lambda: ["--offload-arch=gfx942", "--offload-arch=gfx950"], + ) with warnings.catch_warnings(): warnings.simplefilter("ignore") From 19425784b47695114c6f37c3520c4347015e7835 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 14:23:08 -0400 Subject: [PATCH 08/16] fix(rocm): route explicit arch lists through the resolver too, and warn once Addresses all four suppressed comments in Copilot review 4985982606 on #289. Three continue findings the previous commits fixed only partway. - The validators consulted resolve_target_archs() only when arch_list was None, so an explicit caller bypassed the canonicalization added earlier in this branch: validate_rocm_arch("gfx950:sramecc+") reached the compatibility matrix with the qualifier attached and failed on a value the resolver handles, and validate_flashinfer_rocm_arch("gfx942;gfx950") still saw one invalid token. Both now call resolve_target_archs(arch_list) unconditionally, which is what "single source of truth" was supposed to mean -- the previous form routed only half the callers through it. - Caching rocminfo_gpu_agents stopped the repeated subprocess but not the repeated warning, which lives past the cached call. gen_jit_spec() validates per module, so one JIT generation printed the GPU-less fallback warning once per operation. Measured: 5 resolves emitted 5 warnings, now 1. Keyed on the message so a genuine change of cause -- unset becoming set-but-empty -- is still reported. - test_environment_gets_the_validated_list_not_the_resolved_one also depended on the host's ROCm version: the 6.3/6.4 compatibility path supports neither gfx900 nor gfx950, so CompilationContext would raise before the republished value was checked. Now stubs get_system_rocm_version alongside the PyTorch flags. That is the third environment dependency found in these tests; the pattern is that anything reaching CompilationContext needs both stubs. --- flashinfer/hip_utils.py | 31 ++++++++++++++++++++++++++----- tests/rocm_tests/test_aot_hip.py | 4 ++++ 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 914601b478..4a105c7c7b 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -108,6 +108,20 @@ def resolve_target_archs(arch_list: str = None) -> str: if from_env is None else f"FLASHINFER_ROCM_ARCH_LIST={from_env!r} names no architecture" ) + _warn_gpuless_fallback(how, fallback) + return fallback + + +@functools.cache +def _warn_gpuless_fallback(how: str, fallback: str) -> None: + """Emit the GPU-less fallback warning once per distinct message. + + Caching ``rocminfo_gpu_agents`` stopped the repeated subprocess but not + this: ``gen_jit_spec`` calls ``check_rocm_arch()`` for every module, so a + single JIT generation would otherwise print this once per operation. The + cache key is the message itself, so a genuine change of cause -- unset + becoming set-but-empty -- still gets reported. + """ logger.warning( "No supported AMD GPU detected and %s; building for every supported " "architecture (%s). This is slower than targeting one. Set " @@ -115,7 +129,6 @@ def resolve_target_archs(arch_list: str = None) -> str: how, fallback, ) - return fallback def get_rocm_home(): @@ -337,8 +350,12 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: } # Get architecture list from parameter, env var, or default - if arch_list is None: - arch_list = resolve_target_archs() + # Always through the resolver, even when the caller supplied a list: it is + # what canonicalizes qualifiers and ';' separators, and routing only the + # None case through it left an explicit caller bypassing that -- + # validate_rocm_arch("gfx950:sramecc+") reached the compatibility matrix + # with the qualifier attached and failed on a list the resolver handles. + arch_list = resolve_target_archs(arch_list) # Validate system has ROCm installed system_rocm_version = get_system_rocm_version() @@ -421,8 +438,12 @@ def validate_flashinfer_rocm_arch( """ # Get architecture list from parameter, env var, or default - if arch_list is None: - arch_list = resolve_target_archs() + # Always through the resolver, even when the caller supplied a list: it is + # what canonicalizes qualifiers and ';' separators, and routing only the + # None case through it left an explicit caller bypassing that -- + # validate_rocm_arch("gfx950:sramecc+") reached the compatibility matrix + # with the qualifier attached and failed on a list the resolver handles. + arch_list = resolve_target_archs(arch_list) # Step 1: Validate against system ROCm version (reuse existing logic) validated_arch_list = validate_rocm_arch(arch_list=arch_list, verbose=verbose) diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index 24babb3ca3..f53a1739c3 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -223,6 +223,10 @@ def test_environment_gets_the_validated_list_not_the_resolved_one( "_get_rocm_arch_flags", lambda: ["--offload-arch=gfx942", "--offload-arch=gfx950"], ) + # ...and the ROCm version, for the same reason. The compatibility matrix on + # the 6.3/6.4 path supports neither gfx900 nor gfx950, so on such a host + # CompilationContext raises before the republished value is ever checked. + monkeypatch.setattr("flashinfer.hip_utils.get_system_rocm_version", lambda: "7.1.0") with warnings.catch_warnings(): warnings.simplefilter("ignore") From 4e50a4dd4044f740ddbd4ae6199bd6eadcadbded Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 15:26:42 -0400 Subject: [PATCH 09/16] fix(rocm): make the GPU-less warning test order-independent; correct a comment Found by reviewing my own diff rather than by the review bot, which is where these should have been caught in the first place. - _warn_gpuless_fallback is cached, so the assertion in test_no_device_builds_for_everything_supported is order-dependent: anything earlier in the process emitting the same message leaves caplog empty and the test fails for a reason unrelated to the resolver. It passes today only because exactly one test produces that message. Cleared per test, matching what TestGetSupportedDeviceIndices already does for its caches. - validate_flashinfer_rocm_arch carried a comment pasted verbatim from validate_rocm_arch, naming the wrong function as the example. Rewritten to the reason that actually applies there: the resolve is deliberately repeated so the function is correct standalone, because callers stub validate_rocm_arch and when they do this is the only thing left to turn None into a target. Removing it as "redundant" was tried first and reverted -- test_no_env_no_arg_follows_the_detected_device stubs validate_rocm_arch, so without this call None reaches .split() and raises. The duplication is load bearing. --- flashinfer/hip_utils.py | 10 ++++------ tests/rocm_tests/test_hip_utils.py | 7 +++++++ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 4a105c7c7b..8b56a15322 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -437,12 +437,10 @@ def validate_flashinfer_rocm_arch( RuntimeError: If any validation step fails with clear error message """ - # Get architecture list from parameter, env var, or default - # Always through the resolver, even when the caller supplied a list: it is - # what canonicalizes qualifiers and ';' separators, and routing only the - # None case through it left an explicit caller bypassing that -- - # validate_rocm_arch("gfx950:sramecc+") reached the compatibility matrix - # with the qualifier attached and failed on a list the resolver handles. + # Resolve here as well as in validate_rocm_arch. The second call is + # idempotent, and this keeps the function correct on its own rather than + # borrowing the inner one's behaviour -- callers stub validate_rocm_arch, + # and when they do, this is the only thing left to turn None into a target. arch_list = resolve_target_archs(arch_list) # Step 1: Validate against system ROCm version (reuse existing logic) diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 98b7677112..c3c9568708 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -16,6 +16,7 @@ import pytest from flashinfer.hip_utils import ( + _warn_gpuless_fallback, FLASHINFER_SUPPORTED_ROCM_ARCHS, check_torch_rocm_compatibility, get_available_gpu_count, @@ -243,6 +244,12 @@ def test_no_device_builds_for_everything_supported(self, monkeypatch, caplog): warning tells the operator how to make it cheaper. """ monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + # The warning is emitted through a cached helper so a multi-op JIT run + # does not repeat it per module. That makes this assertion + # order-dependent: anything earlier in the process producing the same + # message would leave caplog empty here and the test would fail for a + # reason unrelated to the resolver. + _warn_gpuless_fallback.cache_clear() with self._agents(), caplog.at_level("WARNING"): assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) assert "FLASHINFER_ROCM_ARCH_LIST" in caplog.text From 484149370b48f5c9a3a4d2c8a4f5f2851f0ba99f Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 16:51:49 -0400 Subject: [PATCH 10/16] fix(rocm): scope architecture detection to the process-visible GPUs Addresses the thread on hip_utils.py:96, which Copilot escalated from a suppressed comment after I deferred it three rounds running. On review the deferral was wrong: the fix is small and needs no torch. rocminfo reports every physical agent and ignores HIP_VISIBLE_DEVICES -- that is exactly why it is safe to call before the HIP runtime starts, and exactly why it is the wrong answer for "what should this process build for". On a host mixing gfx942 and gfx950, a worker pinned to the gfx950 resolved both, then failed validation against a gfx950-only PyTorch; with a fat wheel it built code for cards it cannot address. visible_gpu_agents() filters the probe by HIP_VISIBLE_DEVICES (falling back to CUDA_VISIBLE_DEVICES, HIP winning when both are set, as the runtime does). Indices are rocminfo's enumeration order, which is the HIP runtime's order, and is what tests/conftest.py already writes into the variable -- so interpreting it needs no HIP runtime and this module stays torch-free. Empty means no devices, which is distinct from unset. A UUID form (GPU-...) cannot be mapped to an enumeration index, so the full list is returned rather than guessed: over-building is recoverable, silently targeting the wrong card is not. Two test hazards this created, both found by running under -n rather than only serially, and both invisible in a serial run: - Every test in test_hip_utils.py that patches an agent list would have it filtered by the worker's pinned index. A module-scope autouse fixture clears visibility; the new tests set it themselves afterwards. Failing only on workers whose index exceeds the patched list made this red under -n and green serially. - test_aot_hip.py's failure-path test patched the raw probe, so the same thing happened there. It now patches visible_gpu_agents instead, which is both the intent ("pretend these are visible") and safe -- clearing the variable in that file could unpin the worker before torch initializes, which is the HSA contention conftest.py exists to prevent. 7 tests added. Verified serially and at -n 2 and -n 4. --- flashinfer/hip_utils.py | 46 +++++++++++++++++++++- tests/rocm_tests/test_aot_hip.py | 7 +++- tests/rocm_tests/test_hip_utils.py | 62 ++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 2 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 8b56a15322..192740f58f 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -44,6 +44,50 @@ def _canonical_arch_list(raw: str) -> str: return ",".join(seen) +def visible_gpu_agents() -> tuple: + """``rocminfo_gpu_agents()`` scoped to the GPUs this process can see. + + ``rocminfo`` reports every physical agent and ignores ``HIP_VISIBLE_DEVICES`` + -- documented in ``tests/rocm_tests/conftest.py`` -- which is why it is safe + to call before the HIP runtime starts, but also why it is the wrong answer + for "what should this process build for". On a host mixing gfx942 and + gfx950, a worker pinned to the gfx950 would otherwise resolve both, and then + fail validation against a gfx950-only PyTorch, or build code for cards it + cannot address. + + Indices are rocminfo's enumeration order, which is the HIP runtime's order, + which is what ``tests/conftest.py`` writes into the variable (it pins with an + index from :func:`get_physical_card_device_indices`). So this stays + torch-free: no HIP runtime is needed to interpret it. + + ``HIP_VISIBLE_DEVICES`` wins over ``CUDA_VISIBLE_DEVICES`` when both are set, + matching the HIP runtime. An empty value means no devices, which is distinct + from unset. A value naming devices by UUID (``GPU-...``) cannot be mapped to + an enumeration index here, so the full list is returned rather than guessing + -- over-building is recoverable, silently targeting the wrong card is not. + """ + import os + + agents = rocminfo_gpu_agents() + raw = os.environ.get("HIP_VISIBLE_DEVICES") + if raw is None: + raw = os.environ.get("CUDA_VISIBLE_DEVICES") + if raw is None: + return agents + if not raw.strip(): + return () + + tokens = [t.strip() for t in raw.split(",") if t.strip()] + if not all(t.isdigit() for t in tokens): + logger.debug( + "Not scoping architecture detection: %r is not a list of device " + "indices, so it cannot be mapped to rocminfo's enumeration order.", + raw, + ) + return agents + return tuple(agents[int(t)] for t in tokens if int(t) < len(agents)) + + def resolve_target_archs(arch_list: str = None) -> str: """Return the architectures to build for, as a comma-separated string. @@ -92,7 +136,7 @@ def resolve_target_archs(arch_list: str = None) -> str: detected = sorted( { arch - for arch, _ in rocminfo_gpu_agents() + for arch, _ in visible_gpu_agents() if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS } ) diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index f53a1739c3..e363677d94 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -159,8 +159,13 @@ def test_failed_validation_leaves_the_environment_alone(monkeypatch, tmp_path): import flashinfer.aot_hip as aot_hip monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + # Patch the visibility-aware accessor, not the raw probe: detection is + # scoped by HIP_VISIBLE_DEVICES, which conftest.py pins per xdist worker, so + # a raw-probe patch would have this one-element list filtered by the + # worker's index and fall through to the fat fallback on every worker but + # gw0 -- green serially, red under -n. monkeypatch.setattr( - "flashinfer.hip_utils.rocminfo_gpu_agents", + "flashinfer.hip_utils.visible_gpu_agents", lambda: (("gfx950", ""),), ) diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index c3c9568708..0e53b2bd30 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -31,6 +31,20 @@ ) +@pytest.fixture(autouse=True) +def _unpinned_devices(monkeypatch): + """Clear device visibility for every test in this module. + + Detection is scoped by HIP_VISIBLE_DEVICES, and tests/conftest.py pins it + per xdist worker. Any test that patches an agent list would otherwise have + that list silently filtered by the worker's index -- and because the index + differs per worker, it fails only under -n, only on some workers. Tests that + care about visibility set the variable themselves, after this. + """ + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + # get_rocm_home class TestGetRocmHome: def test_rocm_path_env_var(self, monkeypatch): @@ -219,6 +233,54 @@ def test_a_value_that_normalizes_away_falls_through(self, monkeypatch, raw): with self._agents("gfx950"): assert resolve_target_archs() == "gfx950" + @pytest.mark.parametrize( + "visible,expected", + [ + # The case that motivated this: a mixed host with the worker pinned + # to the gfx950. Resolving both would fail validation against a + # gfx950-only PyTorch and build code for a card this process cannot + # address. + ("1", "gfx950"), + ("0", "gfx942"), + ("0,1", "gfx942,gfx950"), + # Out-of-range indices are skipped rather than raising. + ("1,7", "gfx950"), + ], + ) + def test_detection_is_scoped_to_visible_devices( + self, monkeypatch, visible, expected + ): + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", visible) + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == expected + + def test_hip_visible_devices_wins_over_cuda(self, monkeypatch): + """Matches the HIP runtime, which ignores CUDA_VISIBLE_DEVICES when both + are set.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == "gfx950" + + def test_uuid_visibility_is_not_guessed(self, monkeypatch): + """A UUID form cannot be mapped to rocminfo's enumeration order. + Over-building is recoverable; targeting the wrong card silently is not. + """ + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "GPU-a1b2c3d4") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == "gfx942,gfx950" + + def test_empty_visibility_means_no_devices(self, monkeypatch): + """Empty is distinct from unset: every GPU is hidden, so there is + nothing to detect and the fat fallback is the honest answer.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + def test_detects_the_running_device(self, monkeypatch): monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) with self._agents("gfx950"): From c3c9fb1b9a0a485c37840dbc9a3e5e71a3222fbb Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 17:05:05 -0400 Subject: [PATCH 11/16] fix(rocm): let the PyTorch check filter instead of failing the whole build Step 3 raised if any target was absent from the wheel, while steps 1 and 2 already drop what they cannot serve and continue. That made this PR's GPU-less fallback unusable: it returns every supported architecture, so an arch-specific PyTorch (gfx942-only) failed the entire build over the gfx950 half rather than building the half it can. Verified: gfx942,gfx950 against a gfx942-only wheel now yields {'gfx942'} with a warning, where it previously raised. Still raises when nothing survives, so an outright wrong target is unchanged. --- flashinfer/hip_utils.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 192740f58f..1dd4db7f15 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -526,6 +526,28 @@ def validate_flashinfer_rocm_arch( missing_in_pytorch = [ flag for flag in arch_flags if flag not in pytorch_arch_flags ] + if missing_in_pytorch: + # Filter, then raise only if nothing is left -- matching steps 1 and + # 2, which already drop what they cannot serve and continue. Step 3 + # raising outright made the GPU-less fallback unusable: it returns + # every supported architecture, so an arch-specific PyTorch wheel + # (gfx942-only) failed the whole build over the gfx950 half rather + # than building the half it can. + usable = [f for f in arch_flags if f in pytorch_arch_flags] + if usable: + import warnings + + warnings.warn( + f"PyTorch was not built for " + f"{', '.join(f.removeprefix('--offload-arch=') for f in missing_in_pytorch)}; " + f"building only for " + f"{', '.join(f.removeprefix('--offload-arch=') for f in usable)}.", + UserWarning, + stacklevel=2, + ) + arch_flags = usable + requested_archs = [f.removeprefix("--offload-arch=") for f in usable] + missing_in_pytorch = [] if missing_in_pytorch: raise RuntimeError( f"PyTorch does not support the following architectures: {', '.join(missing_in_pytorch)}.\n" From fdb8f38ad60f433771766b166baf5c9cad59a6da Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 17:45:34 -0400 Subject: [PATCH 12/16] docs+test(rocm): match the resolver docstring to the code, cover the step-3 filter Both found by self-review of the diff rather than by the review bot. - resolve_target_archs' docstring still described step 3 as "the supported GPUs actually present". Since visible_gpu_agents landed it is the supported GPUs *visible to this process*, which is the whole point of that change. - The previous commit changed the PyTorch check from raise-on-any-missing to filter-and-warn without adding a test. The existing raise test still passes because it requests a single arch that is entirely absent, so it never exercises the new branch. Added the partial case, A/B-verified against c3c9fb1b~1 -- without the filter it fails with "Emitted warnings: []". --- flashinfer/hip_utils.py | 3 ++- tests/rocm_tests/test_hip_utils.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 1dd4db7f15..0aee7bca7d 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -95,7 +95,8 @@ def resolve_target_archs(arch_list: str = None) -> str: 1. ``arch_list``, when a caller passes one explicitly. 2. ``FLASHINFER_ROCM_ARCH_LIST``. - 3. The architectures of the supported GPUs actually present. + 3. The architectures of the supported GPUs visible to this process + (``HIP_VISIBLE_DEVICES``-scoped; see :func:`visible_gpu_agents`). 4. Every architecture FlashInfer supports, with a warning. Step 4 replaces a hard-coded ``"gfx942"`` that three call sites reached diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 0e53b2bd30..855ca3d10d 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -509,6 +509,26 @@ def test_pytorch_validation_raises_when_flag_missing(self): arch_list="gfx942", torch_cpp_ext_module=torch_cpp_ext ) + def test_pytorch_validation_filters_when_some_flags_present(self): + """Partial support degrades with a warning; only a total miss raises. + + Steps 1 and 2 already filter-and-continue. Step 3 raising outright made + the GPU-less fallback unusable: it returns every supported architecture, + so an arch-specific wheel failed the whole build over the half it could + not serve. + """ + torch_cpp_ext = MagicMock() + torch_cpp_ext._get_rocm_arch_flags.return_value = ["--offload-arch=gfx942"] + with ( + self._patch_validate_rocm_arch("gfx942,gfx950"), + pytest.warns(UserWarning, match="PyTorch was not built for gfx950"), + ): + flags, arch_set = validate_flashinfer_rocm_arch( + arch_list="gfx942,gfx950", torch_cpp_ext_module=torch_cpp_ext + ) + assert flags == ["--offload-arch=gfx942"] + assert arch_set == {"gfx942"} + def test_reads_arch_from_env_when_none_given(self, monkeypatch): monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942") with self._patch_validate_rocm_arch("gfx942"): From e512b849a9820e376903cd357959532cbdbf8683 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 17:57:54 -0400 Subject: [PATCH 13/16] fix(rocm): narrow only the GPU-less fallback; handle ROCR and the -1 sentinel - The PyTorch filter added in c3c9fb1b applied to every request, so an explicit gfx942,gfx950 against a gfx942-only wheel quietly succeeded as gfx942 instead of failing. That is worse than the bug it fixed: the artifact no longer satisfies the target it claims. resolve_target_archs_with_origin() now reports explicit/env/detected/fallback, and only "fallback" -- a guess this module made with no hardware to look at -- may be narrowed. resolve_target_archs keeps its signature and delegates. - visible_gpu_agents only read HIP_/CUDA_VISIBLE_DEVICES. ROCR_VISIBLE_DEVICES also applies, and applies *beneath* HIP, so the two compose in a way indices cannot express: honoured alone, declined when combined. - "-1" is the no-device sentinel and is not a digit, so it took the unmappable branch and returned every agent -- the exact opposite of what it asks for. The test added in fdb8f38a used an explicit list and would now be wrong; it is replaced by a pair pinning both sides -- fallback narrows with a warning, explicit still raises. --- flashinfer/hip_utils.py | 71 +++++++++++++++++++++++++----- tests/rocm_tests/test_hip_utils.py | 57 +++++++++++++++++++----- 2 files changed, 108 insertions(+), 20 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 0aee7bca7d..84064ebfb5 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -69,15 +69,45 @@ def visible_gpu_agents() -> tuple: import os agents = rocminfo_gpu_agents() - raw = os.environ.get("HIP_VISIBLE_DEVICES") - if raw is None: - raw = os.environ.get("CUDA_VISIBLE_DEVICES") - if raw is None: + # ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set + # alongside one of the others the two compose and the composition cannot be + # reconstructed from indices alone. Refuse to scope in that case rather than + # apply the wrong one -- over-building is recoverable. + present = { + name: os.environ[name] + for name in ( + "HIP_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + ) + if name in os.environ + } + if not present: + return agents + if "ROCR_VISIBLE_DEVICES" in present and len(present) > 1: + logger.debug( + "Not scoping architecture detection: ROCR_VISIBLE_DEVICES composes " + "with %s and the result cannot be mapped to rocminfo's enumeration " + "order.", + ", ".join(n for n in present if n != "ROCR_VISIBLE_DEVICES"), + ) return agents + # HIP wins over CUDA when both are set, as the HIP runtime does. + raw = ( + present.get("HIP_VISIBLE_DEVICES") + or present.get("ROCR_VISIBLE_DEVICES") + or present.get("CUDA_VISIBLE_DEVICES") + or "" + ) if not raw.strip(): return () tokens = [t.strip() for t in raw.split(",") if t.strip()] + # -1 is the documented "no devices" sentinel. It is not a digit, so without + # this it would fall into the unmappable branch below and return *every* + # agent -- the exact opposite of what it asks for. + if any(t.lstrip("+-").isdigit() and int(t) < 0 for t in tokens): + return () if not all(t.isdigit() for t in tokens): logger.debug( "Not scoping architecture detection: %r is not a list of device " @@ -117,6 +147,20 @@ def resolve_target_archs(arch_list: str = None) -> str: before the HIP runtime starts, and keeps this module importable without torch -- see ``tests/rocm_tests/test_arch_caps_hip.py``. """ + return resolve_target_archs_with_origin(arch_list)[0] + + +def resolve_target_archs_with_origin(arch_list: str = None) -> tuple: + """:func:`resolve_target_archs`, plus where the answer came from. + + Origin is one of ``"explicit"``, ``"env"``, ``"detected"``, ``"fallback"``. + Callers need it because the four are not equally authoritative: the + ``"fallback"`` list is a guess this module made when it could not see any + GPU, so narrowing it to what PyTorch actually ships is a courtesy. Narrowing + an ``"explicit"`` request is not -- the operator asked for those + architectures and silently delivering fewer is how an artifact ends up not + satisfying the target it was built for. + """ import os # Canonicalize the two operator-supplied paths. A value that normalizes away @@ -126,13 +170,13 @@ def resolve_target_archs(arch_list: str = None) -> str: if arch_list: canonical = _canonical_arch_list(arch_list) if canonical: - return canonical + return canonical, "explicit" from_env = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") if from_env: canonical = _canonical_arch_list(from_env) if canonical: - return canonical + return canonical, "env" detected = sorted( { @@ -142,7 +186,7 @@ def resolve_target_archs(arch_list: str = None) -> str: } ) if detected: - return ",".join(detected) + return ",".join(detected), "detected" fallback = ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) # Distinguish the two ways of getting here. Telling an operator who *did* @@ -154,7 +198,7 @@ def resolve_target_archs(arch_list: str = None) -> str: else f"FLASHINFER_ROCM_ARCH_LIST={from_env!r} names no architecture" ) _warn_gpuless_fallback(how, fallback) - return fallback + return fallback, "fallback" @functools.cache @@ -486,7 +530,7 @@ def validate_flashinfer_rocm_arch( # idempotent, and this keeps the function correct on its own rather than # borrowing the inner one's behaviour -- callers stub validate_rocm_arch, # and when they do, this is the only thing left to turn None into a target. - arch_list = resolve_target_archs(arch_list) + arch_list, origin = resolve_target_archs_with_origin(arch_list) # Step 1: Validate against system ROCm version (reuse existing logic) validated_arch_list = validate_rocm_arch(arch_list=arch_list, verbose=verbose) @@ -534,8 +578,15 @@ def validate_flashinfer_rocm_arch( # every supported architecture, so an arch-specific PyTorch wheel # (gfx942-only) failed the whole build over the gfx950 half rather # than building the half it can. + # Only the GPU-less fallback may be narrowed. It is a guess this + # module made with no hardware to look at, so trimming it to what + # PyTorch ships is strictly better than failing. An explicit list -- + # from the caller, the environment, or detected hardware -- is a + # stated requirement, and quietly delivering fewer architectures + # than were asked for produces an artifact that does not satisfy the + # target it claims to be built for. usable = [f for f in arch_flags if f in pytorch_arch_flags] - if usable: + if usable and origin == "fallback": import warnings warnings.warn( diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 855ca3d10d..0a05f14131 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -43,6 +43,7 @@ def _unpinned_devices(monkeypatch): """ monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising=False) # get_rocm_home @@ -264,6 +265,31 @@ def test_hip_visible_devices_wins_over_cuda(self, monkeypatch): with self._agents("gfx942", "gfx950"): assert resolve_target_archs() == "gfx950" + @pytest.mark.parametrize("var", ["HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"]) + def test_minus_one_means_no_devices(self, monkeypatch, var): + """-1 is the documented no-device sentinel. It is not a digit, so + without special handling it took the unmappable branch and returned + every agent -- the opposite of what it asks for.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv(var, "-1") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + + def test_rocr_visible_devices_is_honoured_alone(self, monkeypatch): + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == "gfx950" + + def test_rocr_combined_with_hip_is_not_composed(self, monkeypatch): + """ROCr applies beneath HIP, so the two compose and the result cannot be + reconstructed from indices. Decline to scope rather than apply one.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == "gfx942,gfx950" + def test_uuid_visibility_is_not_guessed(self, monkeypatch): """A UUID form cannot be mapped to rocminfo's enumeration order. Over-building is recoverable; targeting the wrong card silently is not. @@ -509,14 +535,12 @@ def test_pytorch_validation_raises_when_flag_missing(self): arch_list="gfx942", torch_cpp_ext_module=torch_cpp_ext ) - def test_pytorch_validation_filters_when_some_flags_present(self): - """Partial support degrades with a warning; only a total miss raises. - - Steps 1 and 2 already filter-and-continue. Step 3 raising outright made - the GPU-less fallback unusable: it returns every supported architecture, - so an arch-specific wheel failed the whole build over the half it could - not serve. - """ + def test_pytorch_filtering_applies_only_to_the_gpuless_fallback(self, monkeypatch): + """The fallback is a guess made with no hardware in sight, so trimming + it to what PyTorch ships beats failing.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setattr("flashinfer.hip_utils.visible_gpu_agents", lambda: ()) + _warn_gpuless_fallback.cache_clear() torch_cpp_ext = MagicMock() torch_cpp_ext._get_rocm_arch_flags.return_value = ["--offload-arch=gfx942"] with ( @@ -524,11 +548,24 @@ def test_pytorch_validation_filters_when_some_flags_present(self): pytest.warns(UserWarning, match="PyTorch was not built for gfx950"), ): flags, arch_set = validate_flashinfer_rocm_arch( - arch_list="gfx942,gfx950", torch_cpp_ext_module=torch_cpp_ext + arch_list=None, torch_cpp_ext_module=torch_cpp_ext ) - assert flags == ["--offload-arch=gfx942"] assert arch_set == {"gfx942"} + def test_explicit_request_still_fails_when_pytorch_lacks_one(self): + """An explicit list is a stated requirement, not a suggestion. Silently + building fewer architectures yields an artifact that does not satisfy + the target it claims.""" + torch_cpp_ext = MagicMock() + torch_cpp_ext._get_rocm_arch_flags.return_value = ["--offload-arch=gfx942"] + with ( + self._patch_validate_rocm_arch("gfx942,gfx950"), + pytest.raises(RuntimeError, match="PyTorch does not support"), + ): + validate_flashinfer_rocm_arch( + arch_list="gfx942,gfx950", torch_cpp_ext_module=torch_cpp_ext + ) + def test_reads_arch_from_env_when_none_given(self, monkeypatch): monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942") with self._patch_validate_rocm_arch("gfx942"): From a421230e6eed5f5687bb6c8ef6dcab3a76830db4 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 18:14:23 -0400 Subject: [PATCH 14/16] fix(rocm): keep resolution provenance through the callers; fix visibility precedence The provenance added in e512b849 was discarded by both callers, so the fallback-only narrowing it gates never fired where it was needed. CompilationContext resolved to a string and passed it back in, which re-entered the resolver as "explicit"; aot_hip published to the environment first, which made it "env". A GPU-less host with an arch-specific wheel therefore still failed on the half it could not build. Verified end to end: that case now yields gfx942 with a warning, where it raised before. - CompilationContext passes arch_list=None and lets validation resolve, so the origin survives. It logs the validated set instead of the pre-resolved string. - aot_hip validates first and publishes after. Resolving twice is cheap now -- the probe is cached and the GPU-less warning is emitted once per message -- and publishing only on success also removes the need to restore on failure. Two precedence bugs in visible_gpu_agents, both from using `or` over values that are legitimately falsy: - HIP_VISIBLE_DEVICES="" means no devices, but empty is falsy, so the chain fell through to CUDA_VISIBLE_DEVICES and resolved GPUs the operator had hidden. Selection is now by key, in strict precedence order. - HIP_VISIBLE_DEVICES=-1 combined with any ROCR value hit the composition early return and got every agent. "No devices" is unambiguous whatever else is set, so it is now decided before the composition check. test_failed_validation_leaves_the_environment_alone asserted the old publish-before-validate order; it now asserts the reverse, which is the property that keeps provenance intact. --- flashinfer/aot_hip.py | 43 ++++++++++----------------- flashinfer/compilation_context_hip.py | 15 ++++++---- flashinfer/hip_utils.py | 34 +++++++++++---------- tests/rocm_tests/test_aot_hip.py | 9 +++--- 4 files changed, 50 insertions(+), 51 deletions(-) diff --git a/flashinfer/aot_hip.py b/flashinfer/aot_hip.py index 931f271de0..626f44ed1c 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -222,38 +222,27 @@ def compile_and_package_modules( # boundary. That is a wider change than this one; leaving the lifetime # unchanged here keeps this commit to the resolution bug it is fixing. from .compilation_context_hip import CompilationContext - from .hip_utils import resolve_target_archs - # Publish first so CompilationContext() -- which resolves through - # resolve_target_archs() itself -- validates exactly this list rather than - # re-deriving one. rocminfo_gpu_agents() is cached now, so this is no longer - # about avoiding a second subprocess; it is about the two agreeing by - # construction instead of by coincidence, which is the whole point of the PR. + # Validate first, publish after. Publishing before would make + # CompilationContext resolve from the environment and classify the GPU-less + # fallback as an "env" request, which forfeits the fallback-only narrowing + # and fails an arch-specific wheel on the half it cannot build. Resolving + # twice is cheap now: rocminfo_gpu_agents is cached and the GPU-less warning + # is emitted once per message. # - # Restore on failure so a build that dies here does not leave the variable - # pointing somewhere for whatever runs next in-process. - rocm_arch_list = resolve_target_archs() - previous_arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") - os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list - try: - context = CompilationContext() # raises if nothing in the list is usable - except BaseException: - if previous_arch_list is None: - os.environ.pop("FLASHINFER_ROCM_ARCH_LIST", None) - else: - os.environ["FLASHINFER_ROCM_ARCH_LIST"] = previous_arch_list - raise - - # Republish what will *actually* be compiled. Validation drops architectures + # Publishing only on success also means a failed build cannot leave + # FLASHINFER_ROCM_ARCH_LIST pointing somewhere for whatever runs next. + context = CompilationContext() # raises if nothing in the list is usable + + # Publish what will *actually* be compiled. Validation drops architectures # it cannot serve with a warning rather than raising, provided at least one - # survives -- "gfx900,gfx942" validates down to "gfx942" -- so the resolved - # list can be a strict superset of the built one. Leaving the superset in - # the environment is the exact failure this PR exists to remove: the AITER - # shim resolves its own target from this variable, so it would build for an - # architecture the packaged kernels were not compiled for. + # survives, so the resolved list can be a strict superset of the built one. + # The AITER shim resolves its own target from this variable, so leaving the + # superset would have it build for an architecture the packaged kernels were + # not compiled for -- the exact failure this PR exists to remove. # # Taken from arch_flags rather than TARGET_ROCM_ARCHS because that is a set; - # order is meaningful here, both on the hipcc command line and to AITER. + # order is meaningful, both on the hipcc command line and to AITER. rocm_arch_list = ",".join( flag.removeprefix("--offload-arch=") for flag in context.arch_flags ) diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index ee8cdd74f0..4935a3ce68 100644 --- a/flashinfer/compilation_context_hip.py +++ b/flashinfer/compilation_context_hip.py @@ -50,15 +50,20 @@ def __init__(self): # One resolver for every path that asks "what are we building for", so # this cannot disagree with the validation in hip_utils -- it used to, # returning gfx950 here while validation checked gfx942. - arch_list = hip_utils.resolve_target_archs() - logger.info(f"Target ROCm architectures: {arch_list}") - - # Comprehensive validation (all 3 checks) + # Pass None rather than a pre-resolved string. Resolution carries + # provenance -- explicit / env / detected / fallback -- and handing back + # the resolved string would re-enter the resolver as "explicit", + # discarding it. Only the GPU-less "fallback" may be narrowed to what + # PyTorch actually ships, so losing the origin here turns a buildable + # half into a hard failure on an arch-specific wheel. self.arch_flags, self.TARGET_ROCM_ARCHS = ( hip_utils.validate_flashinfer_rocm_arch( - arch_list=arch_list, torch_cpp_ext_module=torch_cpp_ext, verbose=False + arch_list=None, torch_cpp_ext_module=torch_cpp_ext, verbose=False ) ) + logger.info( + "Target ROCm architectures: %s", ",".join(sorted(self.TARGET_ROCM_ARCHS)) + ) def get_hipcc_flags_list(self) -> list[str]: """ diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 84064ebfb5..af0d079614 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -84,6 +84,25 @@ def visible_gpu_agents() -> tuple: } if not present: return agents + + # Strict precedence, by key. An `or` chain is wrong here: HIP_VISIBLE_DEVICES="" + # and "-1" both mean "no devices" and are falsy/non-digit, so `or` would skip + # past them to CUDA_VISIBLE_DEVICES and resolve GPUs the operator hid. + for name in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + if name in present: + raw = present[name] + break + + # "No devices" is unambiguous whatever else is set, so it is decided before + # the composition check below -- HIP=-1 with ROCR set still means none. + tokens = [t.strip() for t in raw.split(",") if t.strip()] + if not tokens or any(t.lstrip("+-").isdigit() and int(t) < 0 for t in tokens): + return () + + # ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set + # alongside one of the others the two compose and the composition cannot be + # reconstructed from indices alone. Refuse to scope rather than apply the + # wrong one -- over-building is recoverable. if "ROCR_VISIBLE_DEVICES" in present and len(present) > 1: logger.debug( "Not scoping architecture detection: ROCR_VISIBLE_DEVICES composes " @@ -92,22 +111,7 @@ def visible_gpu_agents() -> tuple: ", ".join(n for n in present if n != "ROCR_VISIBLE_DEVICES"), ) return agents - # HIP wins over CUDA when both are set, as the HIP runtime does. - raw = ( - present.get("HIP_VISIBLE_DEVICES") - or present.get("ROCR_VISIBLE_DEVICES") - or present.get("CUDA_VISIBLE_DEVICES") - or "" - ) - if not raw.strip(): - return () - tokens = [t.strip() for t in raw.split(",") if t.strip()] - # -1 is the documented "no devices" sentinel. It is not a digit, so without - # this it would fall into the unmappable branch below and return *every* - # agent -- the exact opposite of what it asks for. - if any(t.lstrip("+-").isdigit() and int(t) < 0 for t in tokens): - return () if not all(t.isdigit() for t in tokens): logger.debug( "Not scoping architecture detection: %r is not a list of device " diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index e363677d94..cf6262f38b 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -195,10 +195,11 @@ def __init__(self): skip_prebuilt=True, ) - # Validation must see the resolved list, not an unset variable, so that the - # list validated and the list published are the same object rather than two - # independent derivations that happen to agree. - assert _Boom.seen == "gfx950" + # Validation must run *before* anything is published: publishing first + # would make CompilationContext resolve from the environment and classify a + # GPU-less fallback as an "env" request, forfeiting the fallback-only + # narrowing. So the variable is still unset when validation runs. + assert _Boom.seen is None assert "FLASHINFER_ROCM_ARCH_LIST" not in os.environ From 5f35c7691d6d8e1a256eb92f8b8de20dbe829812 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 18:36:37 -0400 Subject: [PATCH 15/16] fix(rocm): make visibility parsing total and check every composing layer Both from Copilot review 4987970293, both mine. - `HIP_VISIBLE_DEVICES=--1` crashed. `"--1".lstrip("+-").isdigit()` is True, so it passed the sentinel guard and then raised ValueError inside int(). Parsing is now total via a helper returning None, and anything unparseable takes the existing unmappable-input path. Visibility comes off a launcher command line, so it has to degrade rather than raise. - The no-device check only looked at the precedence-selected variable, so ROCR_VISIBLE_DEVICES=-1 alongside HIP_VISIBLE_DEVICES=0 fell through to the composition branch and returned every agent, despite ROCr having hidden them all. It now checks the selected variable and ROCR whenever ROCR is present, since ROCr composes beneath HIP rather than being overridden by it. CUDA is still only checked when selected, because HIP overrides it outright. 5 tests. The composing-layer one was vacuous on first write -- with a two-agent list both behaviours return gfx942,gfx950 -- and only distinguishes them with a single-agent list, where ignoring ROCR yields gfx942 and honouring it reaches the fat fallback. Verified against a421230e: 3 malformed cases and the composing-layer case all fail there. --- flashinfer/hip_utils.py | 44 ++++++++++++++++++++++-------- tests/rocm_tests/test_hip_utils.py | 23 ++++++++++++++++ 2 files changed, 55 insertions(+), 12 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index af0d079614..1676005a76 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -88,15 +88,34 @@ def visible_gpu_agents() -> tuple: # Strict precedence, by key. An `or` chain is wrong here: HIP_VISIBLE_DEVICES="" # and "-1" both mean "no devices" and are falsy/non-digit, so `or` would skip # past them to CUDA_VISIBLE_DEVICES and resolve GPUs the operator hid. - for name in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): - if name in present: - raw = present[name] - break - - # "No devices" is unambiguous whatever else is set, so it is decided before - # the composition check below -- HIP=-1 with ROCR set still means none. - tokens = [t.strip() for t in raw.split(",") if t.strip()] - if not tokens or any(t.lstrip("+-").isdigit() and int(t) < 0 for t in tokens): + order = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES") + selected = next(name for name in order if name in present) + + def _index(token): + """The token as a device index, or None if it is not one. + + Parsing has to be total: a malformed sign like ``--1`` passes + ``lstrip("+-").isdigit()`` and then explodes in ``int()``. Visibility + input comes from a launcher command line, so it must degrade rather + than raise. + """ + try: + return int(token) + except ValueError: + return None + + def _says_no_devices(value): + toks = [t.strip() for t in value.split(",") if t.strip()] + if not toks: + return True + return any((i := _index(t)) is not None and i < 0 for t in toks) + + # "No devices" is unambiguous, and can come from a layer that only composes: + # ROCR_VISIBLE_DEVICES=-1 hides everything even when HIP_VISIBLE_DEVICES=0 + # wins precedence. CUDA is not checked unless it was selected, because HIP + # overrides it outright rather than composing with it. + layers = {selected} | ({"ROCR_VISIBLE_DEVICES"} & present.keys()) + if any(_says_no_devices(present[name]) for name in layers): return () # ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set @@ -112,14 +131,15 @@ def visible_gpu_agents() -> tuple: ) return agents - if not all(t.isdigit() for t in tokens): + indices = [_index(t.strip()) for t in present[selected].split(",") if t.strip()] + if any(i is None for i in indices): logger.debug( "Not scoping architecture detection: %r is not a list of device " "indices, so it cannot be mapped to rocminfo's enumeration order.", - raw, + present[selected], ) return agents - return tuple(agents[int(t)] for t in tokens if int(t) < len(agents)) + return tuple(agents[i] for i in indices if i < len(agents)) def resolve_target_archs(arch_list: str = None) -> str: diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 0a05f14131..8dbcea55e8 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -281,6 +281,29 @@ def test_rocr_visible_devices_is_honoured_alone(self, monkeypatch): with self._agents("gfx942", "gfx950"): assert resolve_target_archs() == "gfx950" + def test_no_device_from_a_composing_layer_wins(self, monkeypatch): + """ROCR hides everything even when HIP wins precedence: ROCr applies + beneath HIP, so nothing is left for HIP's index to select.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "-1") + # A single gfx942 agent, so the two behaviours are distinguishable: + # treating ROCR=-1 as "no devices" reaches the fat fallback + # (gfx942,gfx950), while ignoring it yields the detected gfx942. With a + # two-agent list both paths return gfx942,gfx950 and the test proves + # nothing. + with self._agents("gfx942"): + assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + + @pytest.mark.parametrize("bad", ["--1", "1x", "0,--2", "+-3"]) + def test_malformed_tokens_degrade_instead_of_raising(self, monkeypatch, bad): + """Visibility comes from a launcher command line. "--1" passes + lstrip("+-").isdigit() and then raised ValueError in int().""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", bad) + with self._agents("gfx942", "gfx950"): + assert resolve_target_archs() == "gfx942,gfx950" + def test_rocr_combined_with_hip_is_not_composed(self, monkeypatch): """ROCr applies beneath HIP, so the two compose and the result cannot be reconstructed from indices. Decline to scope rather than apply one.""" From 14fb0e50470ea195f9b7e05b2abb7b32321f74d2 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Thu, 20 Aug 2026 18:56:30 -0400 Subject: [PATCH 16/16] fix(rocm): leave ROCR scoping to rocminfo, which already applies it My ROCR handling was wrong in concept, not just in detail. rocminfo is an HSA client, so ROCr has already filtered its output; re-applying ROCR_VISIBLE_DEVICES indices on top filtered twice, and the "composition is ambiguous" case I invented to work around it does not exist. Measured on this host: rocminfo -> 1 GPU agent ROCR_VISIBLE_DEVICES=-1 rocminfo -> 0 GPU agents (ROCr scopes it) HIP_VISIBLE_DEVICES=-1 rocminfo -> 1 GPU agent (HIP does not) That asymmetry is the entire reason this function exists, and it also makes HIP compose with ROCr for free: HIP indices index into the ROCr-visible set, which is exactly what rocminfo reports. Only HIP, then CUDA, is applied here now, and the special-casing is gone. Unmappable visibility (a UUID) now reports "nothing known to be visible" rather than the full agent list. Returning everything labelled it `detected`, which the provenance rule forbids narrowing, so an arch-specific wheel became a hard failure. Reporting nothing reaches the fat fallback, which may be narrowed -- and warns, since a fat build is not what the operator asked for. Two self-inflicted test failures on the way, both caught before pushing: a stale test from the old model survived my replacement block, and removing it took the parametrize decorator off the neighbouring test. --- flashinfer/hip_utils.py | 107 +++++++++-------------------- tests/rocm_tests/test_hip_utils.py | 60 ++++++++-------- 2 files changed, 63 insertions(+), 104 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 1676005a76..5d40bb546d 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -47,98 +47,59 @@ def _canonical_arch_list(raw: str) -> str: def visible_gpu_agents() -> tuple: """``rocminfo_gpu_agents()`` scoped to the GPUs this process can see. - ``rocminfo`` reports every physical agent and ignores ``HIP_VISIBLE_DEVICES`` - -- documented in ``tests/rocm_tests/conftest.py`` -- which is why it is safe - to call before the HIP runtime starts, but also why it is the wrong answer - for "what should this process build for". On a host mixing gfx942 and - gfx950, a worker pinned to the gfx950 would otherwise resolve both, and then - fail validation against a gfx950-only PyTorch, or build code for cards it - cannot address. + Only ``HIP_VISIBLE_DEVICES`` (then ``CUDA_VISIBLE_DEVICES``) is applied here. + ``ROCR_VISIBLE_DEVICES`` deliberately is not: ``rocminfo`` is an HSA client, + so ROCr has *already* filtered its output, and re-applying those indices + would filter twice. Measured on a single-GPU host: + + rocminfo -> 1 GPU agent + ROCR_VISIBLE_DEVICES=-1 rocminfo -> 0 GPU agents (ROCr scopes it) + HIP_VISIBLE_DEVICES=-1 rocminfo -> 1 GPU agent (HIP does not) + + That asymmetry is the whole reason this function exists, and it also means + HIP indices compose correctly with ROCr for free: they index into the + ROCr-visible set, which is exactly what ``rocminfo`` reports. Indices are rocminfo's enumeration order, which is the HIP runtime's order, - which is what ``tests/conftest.py`` writes into the variable (it pins with an - index from :func:`get_physical_card_device_indices`). So this stays + and is what ``tests/conftest.py`` writes into the variable. So this stays torch-free: no HIP runtime is needed to interpret it. - ``HIP_VISIBLE_DEVICES`` wins over ``CUDA_VISIBLE_DEVICES`` when both are set, - matching the HIP runtime. An empty value means no devices, which is distinct - from unset. A value naming devices by UUID (``GPU-...``) cannot be mapped to - an enumeration index here, so the full list is returned rather than guessing - -- over-building is recoverable, silently targeting the wrong card is not. + An empty value or a negative index means no devices. A value naming devices + by UUID (``GPU-...``) cannot be mapped to an enumeration index, so this + reports "nothing known to be visible" and lets the caller fall back to + building for everything -- which, unlike a detected list, may then be + narrowed to what PyTorch actually ships. """ import os agents = rocminfo_gpu_agents() - # ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set - # alongside one of the others the two compose and the composition cannot be - # reconstructed from indices alone. Refuse to scope in that case rather than - # apply the wrong one -- over-building is recoverable. - present = { - name: os.environ[name] - for name in ( - "HIP_VISIBLE_DEVICES", - "CUDA_VISIBLE_DEVICES", - "ROCR_VISIBLE_DEVICES", - ) - if name in os.environ - } - if not present: + for name in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + if name in os.environ: + raw = os.environ[name] + break + else: return agents - # Strict precedence, by key. An `or` chain is wrong here: HIP_VISIBLE_DEVICES="" - # and "-1" both mean "no devices" and are falsy/non-digit, so `or` would skip - # past them to CUDA_VISIBLE_DEVICES and resolve GPUs the operator hid. - order = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES") - selected = next(name for name in order if name in present) - def _index(token): - """The token as a device index, or None if it is not one. - - Parsing has to be total: a malformed sign like ``--1`` passes - ``lstrip("+-").isdigit()`` and then explodes in ``int()``. Visibility - input comes from a launcher command line, so it must degrade rather - than raise. - """ + # Parsing must be total: "--1" passes lstrip("+-").isdigit() and then + # raises inside int(). This value comes off a launcher command line. try: return int(token) except ValueError: return None - def _says_no_devices(value): - toks = [t.strip() for t in value.split(",") if t.strip()] - if not toks: - return True - return any((i := _index(t)) is not None and i < 0 for t in toks) - - # "No devices" is unambiguous, and can come from a layer that only composes: - # ROCR_VISIBLE_DEVICES=-1 hides everything even when HIP_VISIBLE_DEVICES=0 - # wins precedence. CUDA is not checked unless it was selected, because HIP - # overrides it outright rather than composing with it. - layers = {selected} | ({"ROCR_VISIBLE_DEVICES"} & present.keys()) - if any(_says_no_devices(present[name]) for name in layers): + tokens = [t.strip() for t in raw.split(",") if t.strip()] + indices = [_index(t) for t in tokens] + if not tokens or any(i is not None and i < 0 for i in indices): return () - - # ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set - # alongside one of the others the two compose and the composition cannot be - # reconstructed from indices alone. Refuse to scope rather than apply the - # wrong one -- over-building is recoverable. - if "ROCR_VISIBLE_DEVICES" in present and len(present) > 1: - logger.debug( - "Not scoping architecture detection: ROCR_VISIBLE_DEVICES composes " - "with %s and the result cannot be mapped to rocminfo's enumeration " - "order.", - ", ".join(n for n in present if n != "ROCR_VISIBLE_DEVICES"), - ) - return agents - - indices = [_index(t.strip()) for t in present[selected].split(",") if t.strip()] if any(i is None for i in indices): - logger.debug( - "Not scoping architecture detection: %r is not a list of device " - "indices, so it cannot be mapped to rocminfo's enumeration order.", - present[selected], + logger.warning( + "%s=%r does not name device indices, so the target architecture " + "cannot be scoped to it; building for every supported architecture.", + name, + raw, ) - return agents + return () return tuple(agents[i] for i in indices if i < len(agents)) diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 8dbcea55e8..c800dda854 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -25,6 +25,7 @@ get_system_rocm_version_from_hipconfig, is_therock_build, resolve_target_archs, + resolve_target_archs_with_origin, rocminfo_gpu_agents, validate_flashinfer_rocm_arch, validate_rocm_arch, @@ -275,26 +276,6 @@ def test_minus_one_means_no_devices(self, monkeypatch, var): with self._agents("gfx942", "gfx950"): assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) - def test_rocr_visible_devices_is_honoured_alone(self, monkeypatch): - monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) - monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") - with self._agents("gfx942", "gfx950"): - assert resolve_target_archs() == "gfx950" - - def test_no_device_from_a_composing_layer_wins(self, monkeypatch): - """ROCR hides everything even when HIP wins precedence: ROCr applies - beneath HIP, so nothing is left for HIP's index to select.""" - monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) - monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") - monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "-1") - # A single gfx942 agent, so the two behaviours are distinguishable: - # treating ROCR=-1 as "no devices" reaches the fat fallback - # (gfx942,gfx950), while ignoring it yields the detected gfx942. With a - # two-agent list both paths return gfx942,gfx950 and the test proves - # nothing. - with self._agents("gfx942"): - assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) - @pytest.mark.parametrize("bad", ["--1", "1x", "0,--2", "+-3"]) def test_malformed_tokens_degrade_instead_of_raising(self, monkeypatch, bad): """Visibility comes from a launcher command line. "--1" passes @@ -302,25 +283,42 @@ def test_malformed_tokens_degrade_instead_of_raising(self, monkeypatch, bad): monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) monkeypatch.setenv("HIP_VISIBLE_DEVICES", bad) with self._agents("gfx942", "gfx950"): - assert resolve_target_archs() == "gfx942,gfx950" + assert resolve_target_archs() == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) - def test_rocr_combined_with_hip_is_not_composed(self, monkeypatch): - """ROCr applies beneath HIP, so the two compose and the result cannot be - reconstructed from indices. Decline to scope rather than apply one.""" + def test_rocr_is_left_to_rocminfo(self, monkeypatch): + """ROCR_VISIBLE_DEVICES must NOT be applied here: rocminfo is an HSA + client, so ROCr has already filtered its output. Measured -- + `ROCR_VISIBLE_DEVICES=-1 rocminfo` reports 0 GPU agents while + `HIP_VISIBLE_DEVICES=-1 rocminfo` still reports 1. Re-applying the + indices would filter twice; the fixture stands in for output ROCr has + already scoped. + """ monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") - monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + with self._agents("gfx942"): + assert resolve_target_archs() == "gfx942" + + def test_hip_indices_apply_on_top_of_rocr_scoping(self, monkeypatch): + """HIP indexes into the ROCr-visible set, which is what rocminfo + reports -- so the two compose for free and need no special case.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "0,1") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1") with self._agents("gfx942", "gfx950"): - assert resolve_target_archs() == "gfx942,gfx950" + assert resolve_target_archs() == "gfx950" - def test_uuid_visibility_is_not_guessed(self, monkeypatch): - """A UUID form cannot be mapped to rocminfo's enumeration order. - Over-building is recoverable; targeting the wrong card silently is not. - """ + def test_uuid_visibility_falls_back_rather_than_guessing(self, monkeypatch): + """A UUID cannot be mapped to an enumeration index. Reporting "nothing + known to be visible" reaches the fat fallback, which -- unlike a + detected list -- may still be narrowed to what PyTorch ships. Returning + the full list instead would label it `detected` and turn an + arch-specific wheel into a hard failure.""" monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) monkeypatch.setenv("HIP_VISIBLE_DEVICES", "GPU-a1b2c3d4") with self._agents("gfx942", "gfx950"): - assert resolve_target_archs() == "gfx942,gfx950" + archs, origin = resolve_target_archs_with_origin() + assert archs == ",".join(FLASHINFER_SUPPORTED_ROCM_ARCHS) + assert origin == "fallback" def test_empty_visibility_means_no_devices(self, monkeypatch): """Empty is distinct from unset: every GPU is hidden, so there is