fix(rocm): resolve the build architecture in one place, not three

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

flashinfer/aot_hip.py

@@ -208,25 +208,47 @@ 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
 
-    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"]
+    # 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.
+    #
+    # 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, 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, 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}")
 
     # Print summary
     if verbose:

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,23 @@ 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.
+        # 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
             )
         )
-
-    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]:
         """