diff --git a/CLAUDE.md b/CLAUDE.md index e32775a933..4cd5f4ef54 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,7 +50,7 @@ Details: `pr-workflow` skill. | Run tests (fast) | `pytest -n auto --reruns 2 -m "not slow"` | | Run all tests | `pytest -n auto --reruns 2` | | Clear JIT cache | `rm -rf ~/.cache/flashinfer/` | -| Set target arch | `export FLASHINFER_ROCM_ARCH_LIST="gfx942,gfx950"` | +| Set target arch | `export FLASHINFER_ROCM_ARCH_LIST="gfx942,gfx950"` (commas, semicolons or spaces; required when no GPU is visible — the resolver raises rather than guessing) | | Limit parallel build | `export MAX_JOBS=4` | | Verbose JIT output | `export FLASHINFER_JIT_VERBOSE=1` | | Run linting | `pre-commit run -a` | diff --git a/README.md b/README.md index 4ef47c62ec..ab9a9c9d8c 100644 --- a/README.md +++ b/README.md @@ -540,9 +540,15 @@ check_torch_rocm_compatibility() ``` `flashinfer.hip_utils.validate_flashinfer_rocm_arch` is a related -build-time validator used by `setup.py` to cross-check -`FLASHINFER_ROCM_ARCH_LIST` against ROCm and PyTorch — not typically -called from application code. +build-time validator, used by `CompilationContext` to cross-check the +target architectures against ROCm and PyTorch — not typically called +from application code. + +The target itself comes from `flashinfer.hip_utils.resolve_target_archs`: +`FLASHINFER_ROCM_ARCH_LIST` if set (comma-, semicolon- or +whitespace-separated), else the supported GPUs visible to the process. +With neither it raises rather than guessing, so **a build host with no +GPU attached must set `FLASHINFER_ROCM_ARCH_LIST`**. ## Basic Usage diff --git a/amd-flashinfer-jit-cache/README.md b/amd-flashinfer-jit-cache/README.md index c2bb2f47a9..d60b075ae6 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. With neither it raises rather than guessing, so a build host with no GPU attached must set the variable. 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, separated by commas, semicolons, or whitespace. Unset means detect from the visible GPUs; required when none is visible. - `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 5912977f27..f5b8062d31 100644 --- a/flashinfer/aot_hip.py +++ b/flashinfer/aot_hip.py @@ -187,6 +187,19 @@ def compile_and_package_modules( verbose: Whether to print verbose build output skip_prebuilt: Whether to skip pre-built modules """ + # Validate the target before touching anything process-global. Resolution + # can raise now (a GPU-less host with no FLASHINFER_ROCM_ARCH_LIST takes + # that path routinely), and the workspace overrides below outlive the call: + # a failure after them would leave the JIT pointed at this build's + # directory, so a later build or retry in the same process would write + # there. Nothing here needs the workspace, so it costs nothing to go first. + from .compilation_context_hip import CompilationContext + + context = CompilationContext() # raises if nothing in the list is usable + rocm_arch_list = ",".join(sorted(context.TARGET_ROCM_ARCHS)) + if verbose: + print(f"Target ROCm architectures: {rocm_arch_list}") + # Set environment variable (for potential subprocess spawns) os.environ["FLASHINFER_WORKSPACE_BASE"] = str(build_dir) @@ -208,25 +221,11 @@ def compile_and_package_modules( final_config.update(config) config = final_config - # ROCm Arch: Ensure env var is set or create/validate using CompilationContext - from .compilation_context_hip import CompilationContext - - 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"] + # Validation used to also write the target back to os.environ, which was the + # only channel to the AITER shim's separate resolver. Both now share + # hip_utils.resolve_target_archs(), so the write is gone -- along with a + # side effect outliving the call, an operator-set value overwritten, and the + # variable left repointed when a later stage raised. # Print summary if verbose: @@ -239,7 +238,7 @@ def compile_and_package_modules( print(" f16_dtype:", config["f16_dtype"]) print(" use_sliding_window:", config["use_sliding_window"]) print(" use_logits_soft_cap:", config["use_logits_soft_cap"]) - print(" FLASHINFER_ROCM_ARCH_LIST:", rocm_arch_list) + print(" target ROCm architectures:", rocm_arch_list) # Generate JIT specs if verbose: diff --git a/flashinfer/compilation_context_hip.py b/flashinfer/compilation_context_hip.py index 28eb694f76..332d1c3981 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,40 +47,20 @@ 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}") - - # Comprehensive validation (all 3 checks) + # 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. Constructed + # lazily (jit/core.py's _LazyCompilationContext), so "once" means once + # per process that actually compiles, at the point the environment has + # settled -- not at import. 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 ) ) - - 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" + 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 6a91f0b76b..71231a3f5e 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -3,15 +3,225 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import logging +import re # 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"] +# ',', ';' and whitespace all separate architectures. ';' because +# jit/aiter_source.py documents it for this same variable; whitespace because it +# is what PYTORCH_ROCM_ARCH uses, and CLAUDE.md's own example quotes the value. +_ARCH_SEPARATORS = re.compile(r"[;,\s]+") + + +class TargetArchUnresolved(RuntimeError): + """No target architecture could be determined, and none was supplied. + + Its own type so a caller that can carry on -- a packaging step building for + a named list -- can catch it without swallowing the validation failures + raised below. + """ + + +def _canonical_arch_list(raw: str) -> str: + """``"gfx950:sramecc+; gfx942 gfx942"`` -> ``"gfx950,gfx942"``. + + The validators below split on ``","`` only and match tokens verbatim, so an + unnormalized value is a hard failure rather than an untidiness. + + Only *syntax* is normalized: unknown architectures pass through so the + validators can report them. Order is first-occurrence-wins, which is the + caller's stated preference and what reaches the hipcc command line. + """ + seen = [] + for token in _ARCH_SEPARATORS.split(raw): + arch = normalize_arch(token) + if arch and arch not in seen: + seen.append(arch) + return ",".join(seen) + + +def visible_gpu_agents() -> tuple: + """``rocminfo_gpu_agents()`` scoped to the GPUs this process can see. + + Only ``HIP_VISIBLE_DEVICES`` (then ``CUDA_VISIBLE_DEVICES``) is applied. + ``ROCR_VISIBLE_DEVICES`` deliberately is not: rocminfo is an HSA client, so + ROCr has already filtered its output. Measured on a single-GPU host -- + ``ROCR_VISIBLE_DEVICES=-1 rocminfo`` reports 0 agents, ``HIP_VISIBLE_DEVICES=-1 + rocminfo`` still reports 1. HIP indices therefore compose with ROCr for free. + + Selectors may be indices or ``GPU-`` names; HIP accepts both, and + rocminfo reports the UUID per agent so both map to the same enumeration. + + Parsing **truncates** at the first selector that resolves to nothing, + matching HIP: ``HIP_VISIBLE_DEVICES=0,7`` on a two-GPU host yields one + device. That covers ``-1`` and launcher junk without a case each. + + Indices are rocminfo's enumeration order (the HIP runtime's), so this stays + torch-free -- ``tests/conftest.py`` calls it before torch loads. + """ + import os + + agents = rocminfo_gpu_agents() + for name in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + if name in os.environ: + raw = os.environ[name] + break + else: + return agents + + visible = [] + for token in raw.split(","): + token = token.strip() + try: + index = int(token) + except ValueError: + # Only a non-empty token can name a UUID. An empty one is the + # "hide everything" form, and it would otherwise match the "" + # placeholder stored for an agent whose rocminfo output carries no + # Uuid line -- reporting a GPU as visible for HIP_VISIBLE_DEVICES="". + # The lookup is here rather than above the loop so an all-integer + # value (the common case) never pays for the uuid table. + uuids = _rocminfo_agent_uuids() if token else () + if token not in uuids: + break + index = uuids.index(token) + if not 0 <= index < len(agents): + break + visible.append(agents[index]) + return tuple(visible) + + +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", so the JIT, the AOT + packager and the AITER shim cannot disagree. Resolution order: + + 1. ``arch_list``, when a caller passes one explicitly. + 2. ``FLASHINFER_ROCM_ARCH_LIST``. + 3. The supported GPUs visible to this process (see + :func:`visible_gpu_agents`), then torch as a fallback for hosts with no + rocminfo (see :func:`_torch_detected_archs`). + + With none of those, raise -- no *guessed* step on purpose. A hard-coded arch is + the bug this replaces, and "everything we support" is not reachable by every + consumer: AITER's shim takes exactly one architecture, so a fat list becomes + a gfx942 shim linked beside gfx950 kernels, which segfaults. The CUDA half + of this repo refuses the same way (``flashinfer/aot.py``). + + Detection uses rocminfo rather than ``torch.cuda``, keeping this callable + before the HIP runtime starts and this module importable without torch. + + Raises: + ValueError: ``arch_list`` was given but names no architecture. + TargetArchUnresolved: nothing was given and no supported GPU is visible. + """ + import os + + if arch_list is not None: + # A caller-supplied list that normalizes away is a caller bug, not a + # request for detection: answering with the local hardware would build + # something other than what was asked for. + canonical = _canonical_arch_list(arch_list) + if not canonical: + raise ValueError( + f"arch_list={arch_list!r} names no ROCm architecture. Pass a " + f"comma-separated list such as 'gfx942,gfx950', or None to " + f"resolve from the environment." + ) + return canonical + + from_env = os.environ.get("FLASHINFER_ROCM_ARCH_LIST") + if from_env: + canonical = _canonical_arch_list(from_env) + if canonical: + return canonical + + detected = sorted( + { + arch + for arch, _ in visible_gpu_agents() + if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS + } + ) + if detected: + return ",".join(detected) + + # Unconditional, and torch's own view of visibility is what makes that safe: + # torch.cuda.device_count() honours HIP_VISIBLE_DEVICES, so a deliberate + # hide-all yields no devices here too. Gating this on "rocminfo itself found + # nothing" looked tighter but was worse in both directions -- it skipped the + # fallback whenever visibility parsing merely failed (a GPU- selector + # against a rocminfo that omits Uuid, a working host), and it re-ran the + # probe, which is not cached when empty, doubling a 10 s timeout. + detected = _torch_detected_archs() + if detected: + return ",".join(detected) + + # Telling an operator who *did* set the variable to set it sends them + # looking for something already there. `is None`, not falsiness: an + # explicitly empty value is set, and saying so is the actionable message. + why = ( + "FLASHINFER_ROCM_ARCH_LIST is not set" + if from_env is None + else f"FLASHINFER_ROCM_ARCH_LIST={from_env!r} names no architecture" + ) + raise TargetArchUnresolved( + f"Cannot determine which AMD GPU architecture to build for: no " + f"supported GPU is visible to this process and {why}. Set it to the " + f"architecture you are building for, e.g. " + f"FLASHINFER_ROCM_ARCH_LIST={FLASHINFER_SUPPORTED_ROCM_ARCHS[0]} " + f"(supported: {', '.join(FLASHINFER_SUPPORTED_ROCM_ARCHS)}). If a GPU " + f"is present, check HIP_VISIBLE_DEVICES and that rocminfo runs." + ) + + +def _torch_detected_archs() -> list: + """Supported architectures torch can see, as a last resort before raising. + + rocminfo ships in its own ROCm package, so it can be missing (or merely off + ``PATH``) on a host whose HIP runtime works. The pre-consolidation + ``CompilationContext._auto_detect_archs`` read ``gcnArchName`` from torch and + built fine there; rocminfo-only detection turned those hosts into a hard + failure. + + This only covers hosts where the *architecture* is undetectable. A host with + no ROCm install at all still fails, one step later: ``validate_rocm_arch`` + raises when ``get_system_rocm_version()`` finds nothing to check the + architecture against. + + rocminfo stays first: it needs no HIP runtime, so it is the only probe safe + on the pre-``HIP_VISIBLE_DEVICES`` path in tests/conftest.py. torch is + imported inside the function to keep this module torch-free at import. + """ + try: + import torch + + return sorted( + { + arch + for i in range(torch.cuda.device_count()) + if ( + arch := normalize_arch( + torch.cuda.get_device_properties(i).gcnArchName + ) + ) + in FLASHINFER_SUPPORTED_ROCM_ARCHS + } + ) + except Exception: + logger.debug("torch device probe failed", exc_info=True) + return [] + def get_rocm_home(): """ @@ -190,7 +400,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 +409,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 @@ -233,8 +442,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 = os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942") + # 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() @@ -315,11 +528,12 @@ 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") + # 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) validated_arch_list = validate_rocm_arch(arch_list=arch_list, verbose=verbose) @@ -361,6 +575,10 @@ def validate_flashinfer_rocm_arch( flag for flag in arch_flags if flag not in pytorch_arch_flags ] if missing_in_pytorch: + # Raise rather than narrow, unlike steps 1 and 2: every list that + # reaches here is a stated requirement, since there is no guessed + # fallback left to trim. Building fewer architectures than were + # asked for yields an artifact that does not satisfy its target. raise RuntimeError( f"PyTorch does not support the following architectures: {', '.join(missing_in_pytorch)}.\n" f"PyTorch was compiled with: {', '.join(pytorch_arch_flags)}" @@ -390,6 +608,11 @@ def get_available_gpu_count() -> int: return torch.cuda.device_count() +# ROCR_VISIBLE_DEVICES value -> (agents, uuids), index-aligned. Only ever holds +# non-empty results; see rocminfo_gpu_agents. +_ROCMINFO_CACHE: dict = {} + + def rocminfo_gpu_agents() -> tuple: """ Return ``(arch, marketing_name)`` for each GPU agent rocminfo reports. @@ -403,60 +626,102 @@ 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. + Cached per ``ROCR_VISIBLE_DEVICES`` value -- rocminfo is an HSA client, so + that variable alone changes what it reports. A failed or empty probe is + **not** cached: one 10-second timeout would otherwise pin "no GPUs" for the + process, and ``tests/conftest.py`` gets a single attempt at this path. Returns: tuple[tuple[str, str], ...]: One (arch, marketing_name) pair per GPU agent. Empty if rocminfo is unavailable. """ - import re + import os import subprocess + key = os.environ.get("ROCR_VISIBLE_DEVICES") + cached = _ROCMINFO_CACHE.get(key) + if cached is not None: + return cached[0] + try: result = subprocess.run( ["rocminfo"], capture_output=True, text=True, timeout=10, check=False ) if result.returncode != 0: return () - except (FileNotFoundError, subprocess.TimeoutExpired): + except Exception: + # Everything, not a curated list: PermissionError, OSError from a + # container with no /dev/kfd, UnicodeDecodeError from text=True. All + # mean "could not ask the hardware", and a probe should not propagate. + logger.debug("rocminfo probe failed", exc_info=True) return () agents = [] - name = marketing = None + uuids = [] + name = marketing = uuid = None is_gpu = False def _commit(): if is_gpu: # rocminfo's agent-level "Name:" is normally bare ("gfx942"), but # normalize defensively: an unstripped qualifier would drop every - # GPU from the supported list, and callers would silently fall back - # to the default architecture. + # GPU from the supported list. agents.append((normalize_arch(name or ""), marketing or "")) + uuids.append(uuid or "") for line in result.stdout.splitlines(): s = line.strip() if re.match(r"^Agent \d+", s): _commit() - name = marketing = None + name = marketing = uuid = None is_gpu = False elif s.startswith("Name:") and name is None: name = s.split(":", 1)[1].strip() elif s.startswith("Marketing Name:") and marketing is None: marketing = s.split(":", 1)[1].strip() + elif s.startswith("Uuid:") and uuid is None: + uuid = s.split(":", 1)[1].strip() elif s.startswith("Device Type:") and "GPU" in s: is_gpu = True _commit() # process last agent + if agents: + _ROCMINFO_CACHE[key] = (tuple(agents), tuple(uuids)) return tuple(agents) -@functools.cache +def _rocminfo_agent_uuids() -> tuple: + """``GPU-`` per GPU agent, index-aligned with ``rocminfo_gpu_agents``. + + Kept out of the agent tuple so its documented ``(arch, marketing_name)`` + shape stays put; both come from the one cached probe. ``""`` for an agent + rocminfo reports no Uuid for, which never matches a selector. + """ + import os + + rocminfo_gpu_agents() # populate the cache + cached = _ROCMINFO_CACHE.get(os.environ.get("ROCR_VISIBLE_DEVICES")) + return cached[1] if cached else () + + +def clear_rocminfo_cache() -> None: + """Forget every cached hardware probe. For tests that patch ``subprocess``. + + Covers the rocm-smi probe too: tests/conftest.py calls + ``get_physical_card_device_indices()`` in every xdist worker at collection, + so its cache is always warm from real hardware before any test runs. + """ + _ROCMINFO_CACHE.clear() + _primary_indices.cache_clear() + + def get_supported_device_indices() -> tuple: """ Return the indices of AMD GPUs whose architecture is supported by FlashInfer. - The result is cached so rocminfo is invoked at most once per process. + Not cached itself -- it filters ``rocminfo_gpu_agents``, which does the + caching. A successful probe is still paid once; an empty one stays + retryable. Returns: tuple[int, ...]: Device indices of supported GPUs. Empty tuple if none @@ -475,7 +740,6 @@ def get_supported_device_indices() -> tuple: _PRIMARY_VRAM_RATIO = 0.95 -@functools.cache def get_physical_card_device_indices() -> tuple: """ Return one supported device index per physical AMD card. @@ -491,16 +755,20 @@ def get_physical_card_device_indices() -> tuple: helper returns them unchanged. Falls back to the supported-device list if rocm-smi is unavailable or its output cannot be parsed. - Cached per-process; xdist spawns workers as separate processes, so the - cache is per-worker (the underlying rocm-smi call is paid once per - Python process, not once per test). + The rocm-smi probe is cached per-process. The empty case returns before the + cache, so a failed rocminfo is not pinned for conftest.py's one attempt. """ - import json - import subprocess - supported = get_supported_device_indices() if not supported: return () + return _primary_indices(supported) + + +@functools.cache +def _primary_indices(supported: tuple) -> tuple: + """One index per physical card, from rocm-smi's VRAM report.""" + import json + import subprocess try: result = subprocess.run( diff --git a/flashinfer/jit/aiter_source.py b/flashinfer/jit/aiter_source.py index ec3e819ba4..4f98cf989e 100644 --- a/flashinfer/jit/aiter_source.py +++ b/flashinfer/jit/aiter_source.py @@ -20,7 +20,6 @@ import functools import os -import re import shutil from pathlib import Path from typing import List, Optional, Tuple, Union @@ -32,56 +31,18 @@ from .core import JitSpec, logger -_DEFAULT_BUILD_ARCH = "gfx942" - -# An architecture name and nothing else: "gfx942", "gfx950", "gfx90a". Anchored -# so a token that merely starts with "gfx" cannot smuggle in a path separator. -_ARCH_RE = re.compile(r"^gfx[0-9a-f]+$") - - -def _env_arch_list() -> List[str]: - """FLASHINFER_ROCM_ARCH_LIST as a normalized list, accepting ',' or ';'. - - Tokens are checked *after* normalization, and anything that is not an - architecture name is dropped with a warning. Two reachable reasons, both of - which end up naming a directory: - - - A token that is all qualifier and no architecture (``":sramecc+"``, or a - bare ``":"``) is non-empty as written but normalizes to ``""``, which - would become the build architecture and name a cache directory - ``__aiter-``. - - An arbitrary string reaches that same directory name, so ``"../../tmp"`` - escapes the cache root once the tag is joined onto it. - - ``validate_rocm_arch`` already rejects such a token for the main JIT, but it - only *warns and excludes* unless every entry is bad -- so in a mixed list the - bad entry survives to here. - - Empty tokens are dropped silently: a trailing separator is benign, not a typo - worth reporting. - """ - raw = os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "") - archs = [] - for token in re.split(r"[;,]", raw): - arch = normalize_arch(token) - if not arch: - continue - if not _ARCH_RE.match(arch): - logger.warning( - "Ignoring %r in FLASHINFER_ROCM_ARCH_LIST: not a GPU architecture " - "name (expected e.g. 'gfx942').", - token, - ) - continue - archs.append(arch) - return archs - - +@functools.cache def _detected_device_arch() -> Optional[str]: """The running device's architecture, or None if no GPU is visible. torch is imported lazily: this module is imported during JIT setup, and a GPU-less wheel build must not require a working HIP runtime. + + Cached, and that is a correctness requirement rather than a saving. + ``ensure_aiter_lib`` copies the built ``.so`` into the directory this + ultimately names, and ``aiter_jitspec_flags`` builds the ``-L``/``-rpath`` + from a second call; a probe that answered differently between the two would + point the link line at a directory nothing was written to. """ try: import torch @@ -94,50 +55,70 @@ def _detected_device_arch() -> Optional[str]: return None -@functools.lru_cache(maxsize=1) def resolve_aiter_build_arch() -> str: """Return the **single** GPU architecture to build the AITER shim for. - Deliberately one architecture, never a list. AITER's ``get_gfx()`` resolves - to the *last* ``GPU_ARCHS`` entry rather than the running device, so a - multi-arch value makes AITER's Python-level dispatch believe it is on that - architecture no matter what the hardware is -- and it also flips global - compile flags. One entry makes that failure mode unreachable. - - Resolution order: ``FLASHINFER_ROCM_ARCH_LIST`` -> detected device -> - ``gfx942``. When the environment names several architectures, the running - device's is preferred among them; when the environment names architectures - that exclude the running device, the environment still wins (cross-compiling - is legitimate) but the mismatch is reported, because the resulting shim will - fault on this machine. + The target comes from ``hip_utils.resolve_target_archs()``, the resolver the + JIT kernels use, so the shim and the kernels it links beside cannot be built + for different architectures. This module used to parse the env var itself, + which is how they came to disagree. + + Local to AITER is the narrowing to **one** architecture: ``get_gfx()`` reads + the *last* ``GPU_ARCHS`` entry rather than the running device, so a + multi-arch value misleads AITER's dispatch and flips global compile flags. + The running device wins among the candidates; otherwise the first, warned, + since that shim will fault here. + + Raises: + hip_utils.TargetArchUnresolved: propagated -- a wrong single + architecture here segfaults rather than failing to load. + RuntimeError: from validation, when nothing in the list is buildable. + """ + from ..hip_utils import resolve_target_archs + + # The device is part of the key, not just the target list, because the + # selection below prefers the running device -- keying on the list alone + # meant a changed device silently reused the first answer. Both inputs are + # stable per process (_detected_device_arch is cached), so the pair cannot + # flip between the two calls a single AITER module makes. + return _select_build_arch(resolve_target_archs(), _detected_device_arch()) + + +@functools.lru_cache(maxsize=None) +def _select_build_arch(resolved: str, device_arch: Optional[str]) -> str: + """Pick one architecture out of the shared resolver's answer. + + Candidates come from ``validate_flashinfer_rocm_arch``, not from splitting + the resolved string. Both of its filters -- the ROCm compatibility matrix + and the FlashInfer supported list -- *drop* an architecture rather than + raising, so re-deriving either one here is how the shim and the kernels came + to disagree: on ROCm 6.4 with "gfx950,gfx942" the JIT narrowed to gfx942 + while this picked gfx950. Reusing the validator makes that unrepresentable + instead of merely currently-absent. + + Keyed on (target list, running device) rather than ``maxsize=1``, so a + change to either re-selects, and each distinct pair warns once rather than + once per module. """ - env_archs = _env_arch_list() - device_arch = _detected_device_arch() - - if env_archs: - if device_arch and device_arch in env_archs: - return device_arch - if device_arch: - logger.warning( - "FLASHINFER_ROCM_ARCH_LIST=%s does not include this device's " - "architecture (%s); building the AITER shim for %s. It will not " - "run on this GPU.", - ",".join(env_archs), - device_arch, - env_archs[0], - ) - return env_archs[0] + from ..hip_utils import validate_flashinfer_rocm_arch - if device_arch: - return device_arch + # arch_flags, not the arch *set*: order is meaningful here, and a set would + # make the choice below depend on hash ordering. + arch_flags, _ = validate_flashinfer_rocm_arch(arch_list=resolved) + candidates = [f.removeprefix("--offload-arch=") for f in arch_flags] - logger.warning( - "No ROCm device detected and FLASHINFER_ROCM_ARCH_LIST is unset; " - "building the AITER shim for %s. Set FLASHINFER_ROCM_ARCH_LIST to " - "target a different architecture.", - _DEFAULT_BUILD_ARCH, - ) - return _DEFAULT_BUILD_ARCH + if device_arch in candidates: + return device_arch + if device_arch: + logger.warning( + "The resolved ROCm architectures (%s) do not include this device's " + "architecture (%s); building the AITER shim for %s. It will not run " + "on this GPU.", + ",".join(candidates), + device_arch, + candidates[0], + ) + return candidates[0] def _aiter_cache_tag() -> str: @@ -154,11 +135,11 @@ def _aiter_cache_tag() -> str: arch = resolve_aiter_build_arch() # The tag is joined onto the cache root to create a directory, so # "filesystem-safe" above has to be enforced, not just asserted in prose. - # _env_arch_list already rejects anything that is not an architecture name; - # this keeps the guarantee true for the other two sources of `arch` (the - # device probe and the default) and for any future caller. Loud rather than - # silently sanitized: an arch that needs rewriting means the resolver is - # wrong, and a quietly renamed cache directory would hide that. + # _select_build_arch takes its candidates from validate_flashinfer_rocm_arch + # now, which cannot return anything outside FLASHINFER_SUPPORTED_ROCM_ARCHS; + # this is the remaining guard for the device probe and for future callers. + # Loud rather than sanitized: a quietly renamed cache directory would hide + # the resolver bug behind it. if not arch or arch != Path(arch).name or arch.startswith("."): raise ValueError( f"refusing to build a cache directory name from architecture " @@ -173,11 +154,17 @@ def _aiter_cache_tag() -> str: return f"{arch}__aiter-{version}" -@functools.lru_cache(maxsize=1) def _aiter_libs_dir() -> Path: # Keyed by arch + AITER version so a cached lib is never reused across an - # incompatible arch or a changed AITER ABI. - d = jit_env.FLASHINFER_CACHE_DIR / "aiter_libs" / _aiter_cache_tag() + # incompatible arch or a changed AITER ABI. Memoised on the tag rather than + # on nothing: maxsize=1 froze the first directory for the whole process, + # the same staleness refresh_aiter_jitspec undoes one level up. + return _libs_dir_for(_aiter_cache_tag()) + + +@functools.lru_cache(maxsize=None) +def _libs_dir_for(tag: str) -> Path: + d = jit_env.FLASHINFER_CACHE_DIR / "aiter_libs" / tag d.mkdir(parents=True, exist_ok=True) return d @@ -259,6 +246,16 @@ def ensure_aiter_lib(module_name: str) -> Path: # AITER reads these from the environment at build time. from ..hip_utils import get_rocm_home + # Resolve before touching the environment: this can raise, and it used to + # sit between the AITER_SYMBOL_VISIBLE/AITER_JIT_DIR writes and the + # try/finally that restores them. Not reachable today -- _aiter_libs_dir() + # above resolves first, so a raise happens before any write -- but the + # ordering should not depend on that. + # AITER splits GPU_ARCHS on ';' and validates each entry, so a comma-joined + # list reaches it as one unparseable token. A single architecture sidesteps + # that, and is required regardless; see resolve_aiter_build_arch. + gpu_archs = resolve_aiter_build_arch() + prev = { "AITER_SYMBOL_VISIBLE": os.environ.get("AITER_SYMBOL_VISIBLE"), "AITER_JIT_DIR": os.environ.get("AITER_JIT_DIR"), @@ -267,11 +264,7 @@ def ensure_aiter_lib(module_name: str) -> Path: } os.environ["AITER_SYMBOL_VISIBLE"] = "1" os.environ["AITER_JIT_DIR"] = str(aiter_build_dir) - # AITER splits GPU_ARCHS on ';' and validates each entry, so a comma-joined - # list reaches it as one unparseable token. A single architecture sidesteps - # the separator entirely -- and is required regardless; see - # resolve_aiter_build_arch. - os.environ["GPU_ARCHS"] = resolve_aiter_build_arch() + os.environ["GPU_ARCHS"] = gpu_archs os.environ["ROCM_HOME"] = get_rocm_home() built: Optional[Path] = None diff --git a/flashinfer/jit/core.py b/flashinfer/jit/core.py index af94cf6026..82f7b7a9d2 100644 --- a/flashinfer/jit/core.py +++ b/flashinfer/jit/core.py @@ -2,6 +2,7 @@ import functools import logging import os +import threading from contextlib import nullcontext from datetime import datetime from pathlib import Path @@ -135,20 +136,14 @@ def check_rocm_arch(): """ Validate ROCm architecture compatibility for FlashInfer. - Uses centralized validation from hip_utils to ensure: - 1. System ROCm version supports the architectures - 2. FlashInfer has AMD ports for the architectures - 3. PyTorch was compiled with the architectures + Realizing the compilation context *is* the validation -- its constructor + runs all three checks (system ROCm version, FlashInfer AMD port, PyTorch + build). Going through it rather than re-validating means the answer this + approves is the object that later supplies ``--offload-arch``, and makes + the check free after the first of ``gen_jit_spec``'s many calls. """ - import torch.utils.cpp_extension as torch_cpp_ext - from ..hip_utils import validate_flashinfer_rocm_arch - try: - validate_flashinfer_rocm_arch( - arch_list=None, # Uses FLASHINFER_ROCM_ARCH_LIST env or defaults to gfx942 - torch_cpp_ext_module=torch_cpp_ext, - verbose=False, - ) + current_compilation_context.get_target_archs() except RuntimeError as e: raise RuntimeError(f"ROCm architecture validation failed: {e}") from e @@ -160,7 +155,45 @@ def clear_cache_dir(): shutil.rmtree(jit_env.FLASHINFER_JIT_DIR) -current_compilation_context = CompilationContext() +class _LazyCompilationContext: + """``CompilationContext``, constructed on first attribute access. + + Required, not just an optimization: the resolver raises rather than guessing + (``hip_utils.TargetArchUnresolved``), so eager construction made ``import + flashinfer`` itself fail on a GPU-less box with no arch list set. + + It also stops every import paying for architecture detection -- a rocminfo + subprocess plus amd-smi/dpkg/hipconfig, each with a 5-10 s timeout -- and + fixes an ordering bug where an import running before ``tests/conftest.py`` + pinned ``HIP_VISIBLE_DEVICES`` resolved against the unpinned device set. + + Nearly a drop-in: of the 22 accesses across eight modules, all but one sit + inside function bodies, so the deferral holds for ``import flashinfer`` on + both branches. The exception is ``__main__.py``'s module-scope + ``env_variables`` literal, which realizes the context when the CLI module is + imported. That module reads ``TARGET_CUDA_ARCHS`` and shells out to nvcc, so + it does not run on ROCm either way and is left untouched here. + """ + + def __init__(self): + self._ctx = None + self._lock = threading.Lock() + + def __getattr__(self, name): + # Only reached for names not found normally, so _ctx never recurses. + # + # Double-checked locking, because deferring the construction also moved + # it out from behind the import lock: two threads reaching a first JIT + # build together would each build a context, so "resolved once per + # process" would not hold. The fast path stays lock-free. + if self._ctx is None: + with self._lock: + if self._ctx is None: + self._ctx = CompilationContext() + return getattr(self._ctx, name) + + +current_compilation_context = _LazyCompilationContext() @dataclasses.dataclass diff --git a/flashinfer/jit/env.py b/flashinfer/jit/env.py index 3383898806..4a446f8de2 100644 --- a/flashinfer/jit/env.py +++ b/flashinfer/jit/env.py @@ -205,6 +205,36 @@ def _get_aot_dir_hip(): return _package_root / "data" / "aot" + # A bare architecture name and nothing else, so a value off the environment + # cannot contribute a path separator, a "..", or a leading "/". + _BARE_ARCH_RE = re.compile(r"^gfx[0-9a-f]+$") + + def _arch_cache_key(raw: str) -> str: + """``"gfx950:sramecc+;gfx942"`` -> ``"gfx950_gfx942"``, else ``"noarch"``. + + Unusable tokens are *dropped*, not collapsed to ``noarch``: validation + drops them too, so the surviving names are what actually gets built. + Collapsing meant ``gfx942,/etc`` and ``gfx950,/etc`` both keyed + ``noarch`` while building different architectures, which is the binary + reuse this key exists to prevent. + """ + from ..hip_utils import _canonical_arch_list + + archs = [ + a for a in _canonical_arch_list(raw).split(",") if _BARE_ARCH_RE.match(a) + ] + return "_".join(archs) if archs else "noarch" + + class _UnsupportedCurrentDevice(RuntimeError): + """The current device is a real GPU that FlashInfer has no port for. + + Its own type so the bare ``except`` below can swallow every other + failure without swallowing this one. Catching plain ``RuntimeError`` + there also re-raised torch's "No HIP GPUs are available", so ``import + flashinfer`` failed on a GPU-less host -- where "noarch" is the right + cache directory, not an error. + """ + def _get_workspace_dir_name() -> pathlib.Path: try: import torch @@ -234,17 +264,32 @@ def _get_workspace_dir_name() -> pathlib.Path: from ..hip_utils import FLASHINFER_SUPPORTED_ROCM_ARCHS if arch != "noarch" and arch not in FLASHINFER_SUPPORTED_ROCM_ARCHS: - raise RuntimeError( + raise _UnsupportedCurrentDevice( f"torch.cuda.current_device() is device {torch.cuda.current_device()} " f"with unsupported ROCm architecture '{arch}'. " f"Please set the current device to a supported GPU before importing " f"flashinfer (e.g. torch.cuda.set_device()). " f"Supported architectures: {', '.join(FLASHINFER_SUPPORTED_ROCM_ARCHS)}" ) - except RuntimeError: + except _UnsupportedCurrentDevice: raise except Exception: - arch = "noarch" + # No runtime device to name the cache after -- the GPU-less + # cross-compile case, now that the import survives it. Key on the + # requested target instead of a single shared "noarch": that + # directory holds compiled .so files, and JitSpec.build() skips + # write_ninja() when one is already there, so a gfx942 build and a + # later gfx950 build of the same module would silently share the + # first one's binary. Read the raw variable rather than calling the + # resolver: this runs at import, where resolving would both raise on + # a GPU-less host and undo the deferral jit/core.py just added. + # + # Every token must be a bare architecture name. This value names a + # directory that jit/core.py mkdirs at import and clear_cache_dir() + # rmtrees, and an absolute component does not merely escape the + # cache root -- pathlib discards everything to its left, so + # FLASHINFER_ROCM_ARCH_LIST=/etc yields exactly Path("/etc"). + arch = _arch_cache_key(os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "")) # e.g.: $HOME/.cache/flashinfer/0.5.3/gfx942/ return FLASHINFER_CACHE_DIR / flashinfer_version / arch diff --git a/tests/conftest.py b/tests/conftest.py index 97e737067f..4e4f2eec8b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,10 +24,20 @@ from flashinfer.hip_utils import get_physical_card_device_indices _supported = get_physical_card_device_indices() - _gpu_index = ( - _supported[_worker_idx] if _worker_idx < len(_supported) else _worker_idx - ) - os.environ["HIP_VISIBLE_DEVICES"] = str(_gpu_index) + # Wrap when there are more workers than cards. Handing out _worker_idx gave + # gw1 HIP_VISIBLE_DEVICES=1 on a one-card box, which HIP reads as *no* + # devices. Sharing a card is slower and, on CPX, flaky -- but it runs. + # + # When the probe found nothing (no rocminfo on PATH, or it timed out) the + # card count is unknown, so there is no index that can be wrapped: pinning + # to _worker_idx is the same out-of-range bug in a different disguise, and + # a worker with no visible device cannot run at all. Leave the variable + # alone instead. Every worker then shares every card, which is contended + # and, on CPX, flaky -- but degraded beats dead. + if _supported: + os.environ["HIP_VISIBLE_DEVICES"] = str( + _supported[_worker_idx % len(_supported)] + ) import torch from torch.torch_version import TorchVersion diff --git a/tests/rocm_tests/test_aiter_build_arch_hip.py b/tests/rocm_tests/test_aiter_build_arch_hip.py index a75c590bac..bbaeadc856 100644 --- a/tests/rocm_tests/test_aiter_build_arch_hip.py +++ b/tests/rocm_tests/test_aiter_build_arch_hip.py @@ -4,24 +4,50 @@ """Tests for AITER shim build-architecture resolution. -GPU-free by construction: the device probe is monkeypatched, so both CDNA3 and -CDNA4 behaviour is exercised on whichever card happens to be present -- including -the architecture the host does not have. +GPU-free by construction: both device probes are monkeypatched, so CDNA3 and +CDNA4 behaviour are exercised on whichever card is present -- including the +architecture the host does not have. + +"Both" because they answer different questions: hip_utils reads rocminfo to +decide what to *target*, this module reads torch to decide which target is +*running*. See the ``device_arch`` fixture. """ +import re + import pytest +from flashinfer.hip_utils import TargetArchUnresolved from flashinfer.jit import aiter_source @pytest.fixture(autouse=True) -def _clear_resolver_cache(): - """These are lru_cached; each case needs a clean slate.""" - for fn in (aiter_source.resolve_aiter_build_arch, aiter_source._aiter_libs_dir): - fn.cache_clear() +def _clear_resolver_cache(monkeypatch): + """These are lru_cached; each case needs a clean slate. + + Visibility is cleared too: the shared resolver falls through to rocminfo + when the env var is unset, so a case meaning "no target" would otherwise + resolve to whatever card conftest.py pinned this worker to. + """ + for name in ("_select_build_arch", "_libs_dir_for", "_detected_device_arch"): + # getattr: the device_arch fixture replaces _detected_device_arch with a + # plain lambda, which is still in place when this runs at teardown. + getattr(getattr(aiter_source, name), "cache_clear", lambda: None)() + for var in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setattr("flashinfer.hip_utils.visible_gpu_agents", lambda: ()) + # ...and the torch fallback behind it, or a "no target" case resolves to + # whatever card this host actually has. + monkeypatch.setattr("flashinfer.hip_utils._torch_detected_archs", lambda: []) + # Selection now runs the real validator, which probes the system ROCm + # version. Pin it so these cases assert about arch selection rather than + # about whichever ROCm the host carries. + monkeypatch.setattr("flashinfer.hip_utils.get_system_rocm_version", lambda: "7.1.0") yield - for fn in (aiter_source.resolve_aiter_build_arch, aiter_source._aiter_libs_dir): - fn.cache_clear() + for name in ("_select_build_arch", "_libs_dir_for", "_detected_device_arch"): + # getattr: the device_arch fixture replaces _detected_device_arch with a + # plain lambda, which is still in place when this runs at teardown. + getattr(getattr(aiter_source, name), "cache_clear", lambda: None)() @pytest.fixture @@ -43,10 +69,17 @@ def _warn(msg, *args, **kwargs): @pytest.fixture def device_arch(monkeypatch): - """Pretend the machine has a given architecture (or none).""" + """Pretend the machine has a given architecture (or none). + + Both probes are set; patching one describes a machine that cannot exist. + """ def _set(arch): monkeypatch.setattr(aiter_source, "_detected_device_arch", lambda: arch) + monkeypatch.setattr( + "flashinfer.hip_utils.visible_gpu_agents", + lambda: ((arch, ""),) if arch else (), + ) return _set @@ -60,11 +93,25 @@ def test_follows_the_device_when_env_unset(self, monkeypatch, device_arch, arch) device_arch(arch) assert aiter_source.resolve_aiter_build_arch() == arch - def test_defaults_when_no_device_and_no_env(self, monkeypatch, device_arch): - """GPU-less wheel builds are real, so a last-resort default is kept.""" + def test_no_device_and_no_env_refuses_rather_than_defaulting( + self, monkeypatch, device_arch + ): + """The old default was gfx942 whatever the machine was. That guess is + worse here than elsewhere: a gfx942 shim linked beside gfx950 kernels + segfaults rather than failing to load.""" monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) device_arch(None) - assert aiter_source.resolve_aiter_build_arch() == "gfx942" + with pytest.raises(TargetArchUnresolved): + aiter_source.resolve_aiter_build_arch() + + def test_it_shares_the_resolver_the_jit_kernels_use(self, monkeypatch): + """One answer, so the shim and the kernels beside it cannot target + different architectures. This module parsed the env var itself and had + already drifted: "gfx950 gfx942" reached the JIT as two architectures + and this resolver as none.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx950 gfx942") + monkeypatch.setattr(aiter_source, "_detected_device_arch", lambda: None) + assert aiter_source.resolve_aiter_build_arch() == "gfx950" @pytest.mark.parametrize("sep", [",", ";"]) @pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) @@ -120,33 +167,82 @@ def test_qualifier_only_tokens_never_become_the_arch( to '', so filtering the raw token is not enough. Left unfiltered it is returned verbatim -- ':,gfx942' on a gfx950 host - resolved to '' and tagged the cache directory '__aiter-'. The - upstream arch validation warns about such a token and drops it rather - than raising, so it does reach this resolver. + resolved to '' and tagged the cache directory '__aiter-'. """ monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", env) device_arch("gfx950") - assert "" not in aiter_source._env_arch_list() assert aiter_source.resolve_aiter_build_arch() in ("gfx942", "gfx950") - @pytest.mark.parametrize( - "env", ["../../tmp", "..", "garbage", "../../tmp,gfx942", "gfx942,/etc"] - ) - def test_non_architecture_tokens_are_dropped( - self, monkeypatch, device_arch, env, warnings_logged - ): - """The resolved arch names a cache directory, so an arbitrary string + @pytest.mark.parametrize("env", ["../../tmp,gfx942", "gfx942,/etc"]) + def test_non_architecture_tokens_are_dropped(self, monkeypatch, device_arch, env): + """The chosen arch names a cache directory, so an arbitrary string escapes the cache root: '../../tmp,gfx942' created - ``/aiter_libs/../../tmp/pwned__aiter-0.1.10``. + ``/aiter_libs/../../tmp/pwned__aiter-0.1.10``. validate_rocm_arch + only raises when *every* entry is bad, so a mixed list reaches here. + """ + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", env) + device_arch("gfx950") + junk = env.replace("gfx942", "").strip(",") + # match=, so this asserts the junk token is what was reported rather + # than accepting any UserWarning the call happens to emit. + with pytest.warns(UserWarning, match=re.escape(junk)): + assert aiter_source.resolve_aiter_build_arch() == "gfx942" + + def test_an_unsupported_arch_is_not_chosen(self, monkeypatch, device_arch): + """Selection runs the real validator, so its filters cannot be missed. + + ``validate_flashinfer_rocm_arch`` *drops* an arch FlashInfer cannot + serve rather than raising, so with ``gfx90a,gfx942`` the JIT kernels + build for gfx942 -- while a hand-copied filter here picked gfx90a. + Measured on this branch before the fix. + """ + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx90a,gfx942") + device_arch("gfx950") + with pytest.warns(UserWarning, match="gfx90a"): + assert aiter_source.resolve_aiter_build_arch() == "gfx942" + + def test_the_choice_re_selects_when_the_device_changes( + self, monkeypatch, device_arch + ): + """The selection prefers the running device, so the device belongs in + the cache key. - Upstream validate_rocm_arch warns and excludes such a token but only - raises when *every* entry is bad, so a mixed list still reached here. + Keyed on the target list alone, a first call taken before + torch.cuda.set_device() -- AOT codegen, or a worker where device_count() + is still 0 -- picked candidates[0] and froze it for the process, so the + shim was built for the other card. """ + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942,gfx950") + device_arch(None) + assert aiter_source.resolve_aiter_build_arch() == "gfx942" # candidates[0] + device_arch("gfx950") + assert aiter_source.resolve_aiter_build_arch() == "gfx950" + + def test_a_rocm_version_narrowing_is_honoured_too(self, monkeypatch, device_arch): + """The other filter, and the one the first fix missed. + + On ROCm 6.4 the compatibility matrix has no gfx950, so validate_rocm_arch + narrows the JIT kernels to gfx942 while a FlashInfer-supported-list-only + filter kept gfx950 and built the shim for it. Measured before the fix. + """ + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx950,gfx942") + monkeypatch.setattr( + "flashinfer.hip_utils.get_system_rocm_version", lambda: "6.4.0" + ) + device_arch(None) + with pytest.warns(UserWarning, match="gfx950"): + assert aiter_source.resolve_aiter_build_arch() == "gfx942" + + @pytest.mark.parametrize("env", ["../../tmp", "..", "garbage", "gfx90a"]) + def test_a_list_of_nothing_but_junk_raises(self, monkeypatch, device_arch, env): + """Falling back to the device would silently build something other than + what was asked for: the operator sees a working build and never learns + the value was ignored. The message is the validator's, so it is the same + one the JIT kernels produce for the same input.""" monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", env) device_arch("gfx950") - assert all(a.startswith("gfx") for a in aiter_source._env_arch_list()) - assert aiter_source.resolve_aiter_build_arch() in ("gfx942", "gfx950") - assert any("not a GPU architecture" in w for w in warnings_logged) + with pytest.raises(RuntimeError, match="does not support any"): + aiter_source.resolve_aiter_build_arch() class TestCacheTagIsAPathComponent: diff --git a/tests/rocm_tests/test_aot_hip.py b/tests/rocm_tests/test_aot_hip.py index c3873f2078..bd58f42b98 100644 --- a/tests/rocm_tests/test_aot_hip.py +++ b/tests/rocm_tests/test_aot_hip.py @@ -11,18 +11,22 @@ 3. .so files are created and can be loaded """ +import os import shutil import tempfile +import warnings from pathlib import Path import pytest import torch +import torch.utils.cpp_extension as torch_cpp_ext from flashinfer.aot_hip import ( compile_and_package_modules, gen_all_modules, get_default_config, ) +from flashinfer.hip_utils import TargetArchUnresolved # Skip all tests if HIP is not available pytestmark = pytest.mark.skipif( @@ -139,6 +143,85 @@ def test_compile_and_package_minimal(): shutil.rmtree(out_dir, ignore_errors=True) +@pytest.mark.parametrize("preset", [None, "gfx900,gfx950"]) +def test_the_build_does_not_touch_the_environment(monkeypatch, tmp_path, preset): + """An AOT build must leave FLASHINFER_ROCM_ARCH_LIST exactly as it found it. + + It used to publish the validated list there, as the only channel to the + AITER shim's separate resolver. Both now share hip_utils. + + The set case has to be a value the build would *change*: validation narrows + "gfx900,gfx950" to "gfx950". A preset that survives intact makes the old + write a no-op, and the test then passes with the bug present. + """ + if preset is None: + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + else: + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", preset) + # Patch the visibility-aware accessor, not the raw probe: conftest.py pins + # HIP_VISIBLE_DEVICES per worker, which would filter a raw-probe patch -- + # green serially, red under -n. + monkeypatch.setattr( + "flashinfer.hip_utils.visible_gpu_agents", + lambda: (("gfx950", ""),), + ) + # Pin both sides of validation, so this asserts about the write and not + # about whichever torch wheel and ROCm version the host carries. + monkeypatch.setattr("flashinfer.hip_utils.get_system_rocm_version", lambda: "7.1.0") + monkeypatch.setattr( + torch_cpp_ext, + "_get_rocm_arch_flags", + lambda: ["--offload-arch=gfx942", "--offload-arch=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.get("FLASHINFER_ROCM_ARCH_LIST") == preset + + +def test_an_unresolvable_target_fails_before_codegen(monkeypatch, tmp_path): + """A build that cannot succeed should say why up front: gen_all_modules() + walks the whole config matrix, so failing later turns an actionable "set + FLASHINFER_ROCM_ARCH_LIST" into a wait followed by a stack trace.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setattr("flashinfer.hip_utils.visible_gpu_agents", lambda: ()) + # ...and the torch fallback behind it, or this host's real GPU answers. + monkeypatch.setattr("flashinfer.hip_utils._torch_detected_archs", lambda: []) + monkeypatch.setattr( + "flashinfer.aot_hip.gen_all_modules", + lambda *a, **k: pytest.fail("codegen ran despite an unresolvable target"), + ) + + with pytest.raises(TargetArchUnresolved): + 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, + ) + + def test_module_naming_convention(): """Test that generated module names follow expected conventions.""" f16_dtype = [torch.float16] diff --git a/tests/rocm_tests/test_hip_utils.py b/tests/rocm_tests/test_hip_utils.py index 75ea294247..b5c19d9299 100644 --- a/tests/rocm_tests/test_hip_utils.py +++ b/tests/rocm_tests/test_hip_utils.py @@ -17,16 +17,49 @@ from flashinfer.hip_utils import ( FLASHINFER_SUPPORTED_ROCM_ARCHS, + TargetArchUnresolved, check_torch_rocm_compatibility, get_available_gpu_count, get_rocm_home, get_supported_device_indices, get_system_rocm_version_from_hipconfig, is_therock_build, + resolve_target_archs, + visible_gpu_agents, + clear_rocminfo_cache, validate_flashinfer_rocm_arch, validate_rocm_arch, ) +# Bound before any test patches the module attribute, so the probe's own tests +# can reach the real implementation past the autouse stub. +from flashinfer.hip_utils import _torch_detected_archs as _real_torch_probe + + +@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) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising=False) + # The rocminfo probe is cached per ROCR_VISIBLE_DEVICES value. Tests that + # exercise the real one (the CompilationContext agreement test) would + # otherwise leave this machine's agents cached under the unset key, and the + # next test to patch subprocess.run would assert against them instead. + clear_rocminfo_cache() + # Neutralize the torch fallback. It runs only when rocminfo finds nothing, + # which is exactly the state most tests here construct -- so on a real GPU + # host it would answer for them and every "raises when undeterminable" case + # would silently pass for the wrong reason. Tests that want it patch it. + monkeypatch.setattr("flashinfer.hip_utils._torch_detected_archs", lambda: []) + # get_rocm_home class TestGetRocmHome: @@ -137,7 +170,323 @@ 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" + + @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 what AITER's own GPU_ARCHS uses, so an operator moving + # between the two variables types it here too. + ("gfx942;gfx950", "gfx942,gfx950"), + ("gfx942; gfx950", "gfx942,gfx950"), + # Whitespace is PYTORCH_ROCM_ARCH's separator and CLAUDE.md quotes + # the value, so operators type this shape. Splitting on ',;' alone + # left it as one token that matched nothing. + ("gfx942 gfx950", "gfx942,gfx950"), + ("gfx942\tgfx950", "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_an_env_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" + + @pytest.mark.parametrize("raw", ["", ";;", " ", ","]) + def test_an_explicit_argument_that_normalizes_away_raises(self, monkeypatch, raw): + """Unlike the env var, an argument that names nothing is a caller bug: + falling through to detection builds for whatever is attached, which the + call site cannot distinguish from success.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents("gfx950"), pytest.raises(ValueError, match="names no ROCm"): + resolve_target_archs(raw) + + @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"), + # Truncation at the first unusable index, which is what HIP does. + ("1,7", "gfx950"), + ("0,-1,1", "gfx942"), + ("0,junk", "gfx942"), + ], + ) + 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" + + @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, so nothing is detectable.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv(var, "-1") + with self._agents("gfx942", "gfx950"), pytest.raises(TargetArchUnresolved): + resolve_target_archs() + + @pytest.mark.parametrize("bad", ["--1", "1x", "+-3", "GPU-a1b2c3d4"]) + def test_unmappable_visibility_is_not_a_traceback(self, monkeypatch, bad): + """Neither launcher junk nor HIP's legal UUID form maps to an + enumeration index, so nothing is detectable -- but the failure must be + the resolver's own error, not a ValueError escaping int().""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", bad) + with self._agents("gfx942", "gfx950"), pytest.raises(TargetArchUnresolved): + resolve_target_archs() + + 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") + 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() == "gfx950" + + def test_empty_visibility_means_no_devices(self, monkeypatch): + """Empty is distinct from unset: every GPU is hidden, so nothing is + detectable even though rocminfo reports agents.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + with self._agents("gfx942", "gfx950"), pytest.raises(TargetArchUnresolved): + resolve_target_archs() + + def test_hidden_devices_still_yield_to_the_env_var(self, monkeypatch): + """Cross-compiling on a box whose GPUs are hidden must keep working -- + the raise is only for the case where there is nothing to go on.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx942") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "-1") + 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_torch_answers_when_rocminfo_finds_nothing(self, monkeypatch): + """rocminfo ships in its own ROCm package and is absent from slim images + and from a bare `pip install torch --index-url .../rocm` host. + + The pre-consolidation CompilationContext read gcnArchName from torch and + built fine there; rocminfo-only detection made those hosts hard-fail. + """ + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setattr( + "flashinfer.hip_utils._torch_detected_archs", lambda: ["gfx950"] + ) + with self._agents(): + assert resolve_target_archs() == "gfx950" + + def test_rocminfo_is_preferred_over_torch(self, monkeypatch): + """rocminfo needs no HIP runtime, so it stays first -- it is the only + probe safe on conftest.py's pre-HIP_VISIBLE_DEVICES path. + + Asserts the torch probe is never *called*, not just that rocminfo's + answer comes back: the latter passes with the fallback deleted + entirely, so it would prove nothing about the ordering. + """ + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + called = [] + monkeypatch.setattr( + "flashinfer.hip_utils._torch_detected_archs", + lambda: called.append(1) or ["gfx942"], + ) + with self._agents("gfx950"): + assert resolve_target_archs() == "gfx950" + assert called == [], "torch was probed even though rocminfo answered" + + @pytest.mark.parametrize( + "gcn,count,expected", + [ + ("gfx950:sramecc+:xnack-", 1, ["gfx950"]), # qualifiers stripped + ("gfx90a", 1, []), # unsupported arch excluded + ("gfx942", 0, []), # no devices + ], + ) + def test_the_torch_probe_itself(self, gcn, count, expected): + # _real_torch_probe, not the module attribute: the autouse fixture + # stubs the latter, and monkeypatch.undo() would also revert the env + # clearing the rest of this module depends on. + torch_mock = MagicMock() + torch_mock.cuda.device_count.return_value = count + torch_mock.cuda.get_device_properties.return_value.gcnArchName = gcn + with patch.dict("sys.modules", {"torch": torch_mock}): + assert _real_torch_probe() == expected + + def test_the_torch_probe_swallows_a_broken_runtime(self): + """A torch that imports but cannot talk to the driver must degrade to + the raise, not propagate out of a probe.""" + torch_mock = MagicMock() + torch_mock.cuda.device_count.side_effect = RuntimeError("No HIP GPUs") + with patch.dict("sys.modules", {"torch": torch_mock}): + assert _real_torch_probe() == [] + + def test_no_device_and_no_env_refuses_to_guess(self, monkeypatch): + """The one behaviour this design turns on: we do not invent a target. + + A hard-coded arch produced a silently gfx942-only artifact, and a fat + list is unusable by AITER's single-arch shim. The error names the + variable to set, the only thing the operator can act on. + """ + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + with self._agents(), pytest.raises(TargetArchUnresolved) as exc: + resolve_target_archs() + assert "FLASHINFER_ROCM_ARCH_LIST is not set" in str(exc.value) + + @pytest.mark.parametrize("raw", [";;", ""]) + def test_an_env_value_naming_nothing_says_so(self, monkeypatch, raw): + """Telling an operator who *did* set the variable to set it sends them + looking for something already there. Empty counts as set -- falsiness + rather than ``is None`` reported it as unset.""" + monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", raw) + with self._agents(), pytest.raises(TargetArchUnresolved) as exc: + resolve_target_archs() + assert "names no architecture" in str(exc.value) + + 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 + ) + 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 +537,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"): @@ -278,17 +641,53 @@ def test_pytorch_validation_raises_when_flag_missing(self): arch_list="gfx942", torch_cpp_ext_module=torch_cpp_ext ) + def test_a_partial_pytorch_match_is_never_narrowed(self): + """Every list reaching step 3 is now a stated requirement -- explicit, + from the environment, or read off the hardware -- because there is no + guessed fallback left to trim. 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_no_device_and_no_env_propagates_the_resolver_error(self, monkeypatch): + """The validators do not paper over an unresolvable target either.""" + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setattr("flashinfer.hip_utils.visible_gpu_agents", lambda: ()) + with pytest.raises(TargetArchUnresolved): + validate_flashinfer_rocm_arch(arch_list=None) + 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"): 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"): @@ -297,6 +696,80 @@ def test_verbose_prints_message(self, capsys): assert "gfx942" in captured.out +class TestImportIsLazy: + """``import flashinfer`` must not resolve an architecture. + + The resolver raises when it cannot tell, so a context built at import makes + the *import* fail on a GPU-less box -- including for callers that never + compile. It is also an ordering fix: conftest.py pins HIP_VISIBLE_DEVICES + per worker, and an import that resolved first missed it. + + Out-of-process, since a module-scope side effect cannot be re-observed by + reloading. + """ + + def _run(self, body, **env): + import os + import sys + + return subprocess.run( + [sys.executable, "-c", body], + capture_output=True, + text=True, + timeout=300, + env={**os.environ, **env}, + ) + + def test_import_succeeds_with_no_gpu_and_no_env_var(self): + proc = self._run( + "import flashinfer; print('ok')", + HIP_VISIBLE_DEVICES="-1", + ROCR_VISIBLE_DEVICES="-1", + FLASHINFER_ROCM_ARCH_LIST="", + ) + assert proc.returncode == 0, proc.stderr + assert "ok" in proc.stdout + + def test_the_first_build_is_where_it_refuses(self): + """The error surfaces at the point of compiling, not at import. + + The marker is what distinguishes the two: without it, an eager context + fails at the import line and the stderr assertion still passes. + """ + proc = self._run( + "import flashinfer\n" + "print('imported')\n" + "from flashinfer.jit.core import check_rocm_arch\n" + "check_rocm_arch()\n", + HIP_VISIBLE_DEVICES="-1", + ROCR_VISIBLE_DEVICES="-1", + FLASHINFER_ROCM_ARCH_LIST="", + ) + assert "imported" in proc.stdout, proc.stderr + assert proc.returncode != 0 + assert "FLASHINFER_ROCM_ARCH_LIST" in proc.stderr + + def test_import_does_not_shell_out_to_rocminfo(self): + """The detection probes each carry a 5-10 s timeout, and an import that + never compiles should pay none of them. + + FLASHINFER_ROCM_ARCH_LIST is cleared deliberately: with it set (as the + CI container does) the resolver returns before detection, so even an + eager context would not probe and the test would pass either way. + """ + proc = self._run( + "import subprocess\n" + "seen = []\n" + "real = subprocess.run\n" + "subprocess.run = lambda a, *p, **k: (seen.append(a), real(a, *p, **k))[1]\n" + "import flashinfer\n" + "print(seen)\n", + FLASHINFER_ROCM_ARCH_LIST="", + ) + assert proc.returncode == 0, proc.stderr + assert "rocminfo" not in proc.stdout + + # get_available_gpu_count class TestGetAvailableGpuCount: """ @@ -337,19 +810,23 @@ def test_delegates_to_torch_cuda_device_count(self): _ROCMINFO_GPU_AGENT_TEMPLATE = """\ Agent {idx} Name: {name} + Uuid: {uuid} Device Type: GPU """ -def _make_rocminfo_output(*gpu_names, cpu_first=True): +def _make_rocminfo_output(*gpu_names, cpu_first=True, uuids=None): """Build a synthetic rocminfo output string.""" lines = [_ROCMINFO_HEADER] agent_idx = 1 if cpu_first: lines.append(_ROCMINFO_CPU_AGENT.replace("Agent 1", f"Agent {agent_idx}")) agent_idx += 1 - for name in gpu_names: - lines.append(_ROCMINFO_GPU_AGENT_TEMPLATE.format(idx=agent_idx, name=name)) + for i, name in enumerate(gpu_names): + uuid = uuids[i] if uuids else f"GPU-{i:016x}" + lines.append( + _ROCMINFO_GPU_AGENT_TEMPLATE.format(idx=agent_idx, name=name, uuid=uuid) + ) agent_idx += 1 return "".join(lines) @@ -358,10 +835,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): - get_supported_device_indices.cache_clear() + self._clear_caches() + + @staticmethod + def _clear_caches(): + # get_supported_device_indices is a plain filter now, but the probe + # underneath it is cached: without this the first test in the class + # would pin the probe result and the rest would assert against it + # instead of their own patched subprocess output. + clear_rocminfo_cache() def _run_result(self, stdout, returncode=0): result = MagicMock() @@ -432,6 +917,82 @@ def test_returns_tuple_type(self): indices = get_supported_device_indices() assert isinstance(indices, tuple) + @pytest.mark.parametrize( + "visible,expected", + [ + ("GPU-aaa", "gfx942"), + ("GPU-bbb", "gfx950"), + ("GPU-bbb,GPU-aaa", "gfx942,gfx950"), + # An unknown UUID selects nothing, exactly like an out-of-range index. + ("GPU-zzz", None), + ], + ) + def test_uuid_selectors_are_mapped_to_agents(self, monkeypatch, visible, expected): + """HIP accepts ``GPU-`` as well as indices, and rocminfo reports + the Uuid per agent -- so both map to one enumeration. + + Truncating on a UUID instead made a populated host look GPU-less, which + now raises rather than silently producing a fat build. + """ + output = _make_rocminfo_output("gfx942", "gfx950", uuids=["GPU-aaa", "GPU-bbb"]) + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", visible) + with patch("subprocess.run", return_value=self._run_result(output)): + if expected is None: + with pytest.raises(TargetArchUnresolved): + resolve_target_archs() + else: + assert resolve_target_archs() == expected + + def test_empty_visibility_beats_a_uuidless_agent(self, monkeypatch): + """`HIP_VISIBLE_DEVICES=""` must hide everything even when rocminfo + reports no Uuid for an agent. + + The uuid list stores "" for such an agent, and `"" in uuids` matched it + -- so the canonical hide-everything setting reported agent 0 as visible + and the refuse-to-guess property was lost. Only reachable with a + rocminfo that omits Uuid, which the other tests' fixtures never produce. + """ + output = _make_rocminfo_output("gfx942", uuids=[""]) + monkeypatch.delenv("FLASHINFER_ROCM_ARCH_LIST", raising=False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "") + with patch("subprocess.run", return_value=self._run_result(output)): + assert visible_gpu_agents() == () + with pytest.raises(TargetArchUnresolved): + resolve_target_archs() + + def test_a_failed_probe_is_not_cached(self): + """One 10-second timeout must not pin "no GPUs" for the process. + + conftest.py pins each xdist worker through this path before torch + loads, and gets exactly one attempt. + """ + with patch("subprocess.run", side_effect=subprocess.TimeoutExpired("x", 10)): + assert get_supported_device_indices() == () + output = _make_rocminfo_output("gfx942") + with patch("subprocess.run", return_value=self._run_result(output)): + assert get_supported_device_indices() == (0,) + + def test_undecodable_output_is_swallowed(self): + """text=True decodes, so non-UTF-8 output raises UnicodeDecodeError -- + which a curated (OSError, TimeoutExpired) let escape out of a probe.""" + with patch( + "subprocess.run", + side_effect=UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid"), + ): + assert get_supported_device_indices() == () + + def test_the_cache_is_keyed_on_rocr_visibility(self, monkeypatch): + """rocminfo is an HSA client, so ROCR_VISIBLE_DEVICES changes what it + reports; caching across a change would hand back the other scope.""" + two = _make_rocminfo_output("gfx942", "gfx950") + one = _make_rocminfo_output("gfx950") + with patch("subprocess.run", return_value=self._run_result(two)): + assert get_supported_device_indices() == (0, 1) + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "1") + with patch("subprocess.run", return_value=self._run_result(one)): + assert get_supported_device_indices() == (0,) + # check_torch_rocm_compatibility def _make_torch_mock(hip=None): @@ -555,3 +1116,46 @@ def test_constant_is_non_empty(self): def test_all_entries_start_with_gfx(self): for arch in FLASHINFER_SUPPORTED_ROCM_ARCHS: assert arch.startswith("gfx"), f"Unexpected arch: {arch}" + + +class TestWorkspaceDirIsAPathComponent: + """The GPU-less workspace key comes off the environment and names a + directory that jit/core.py mkdirs at import and clear_cache_dir() rmtrees. + + An absolute component does not merely escape the cache root: pathlib + discards everything to its left, so FLASHINFER_ROCM_ARCH_LIST=/etc made + FLASHINFER_WORKSPACE_DIR exactly Path("/etc"). Measured before the fix. + """ + + @pytest.mark.parametrize( + "raw,expected", + [ + ("gfx942", "gfx942"), + ("gfx950:sramecc+;gfx942", "gfx950_gfx942"), + ("gfx942 gfx950", "gfx942_gfx950"), + ("", "noarch"), + ("/etc", "noarch"), + ("../../pwned", "noarch"), + ("a/b", "noarch"), + (".", "noarch"), + # Unusable tokens are dropped, not collapsed to noarch: validation + # drops them too, so gfx942 is what actually gets built. Collapsing + # made "gfx942,/etc" and "gfx950,/etc" share one directory while + # building different architectures. + ("gfx942,/etc", "gfx942"), + ("gfx950,/etc", "gfx950"), + ], + ) + def test_only_bare_arch_names_key_the_directory(self, raw, expected): + from flashinfer.jit.env import _arch_cache_key + + assert _arch_cache_key(raw) == expected + + @pytest.mark.parametrize("raw", ["/etc", "../../pwned", "gfx942,/etc"]) + def test_the_workspace_stays_under_the_cache_root(self, raw, tmp_path): + import flashinfer.jit.env as jit_env + from flashinfer.jit.env import _arch_cache_key + + resolved = (tmp_path / _arch_cache_key(raw)).resolve() + assert tmp_path.resolve() in resolved.parents, resolved + assert jit_env.FLASHINFER_CACHE_DIR # the real root is still a real path