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

[demandal25]

The bug

Three call sites independently answered "which GPU architecture are we compiling for", and on CDNA4 they disagreed. Measured on a gfx950 host with FLASHINFER_ROCM_ARCH_LIST unset:

ACTUAL device                                   gfx950
validate_flashinfer_rocm_arch(arch_list=None)   ['gfx942']   <-- wrong
CompilationContext().TARGET_ROCM_ARCHS          ['gfx950']
resolve_aiter_build_arch()                      gfx950

The JIT validated gfx942 while compiling gfx950. That check exists to catch "your PyTorch was not built for this architecture", so pointed at the wrong architecture it fails both ways: vacuous on a fat PyTorch that carries both (it passes, having checked something you are not using), and a spurious hard failure on an arch-specific build carrying only gfx950 — refusing a setup that works.

The cause was os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942") reached from two functions in hip_utils, plus two more return "gfx942" lines in CompilationContext._auto_detect_archs. gfx942 was the one architecture where the literal happened to be right, which is why CDNA3 never noticed. A unit test pinned the wrong constant (test_defaults_to_gfx942_when_no_env_and_no_arg), which is how it survived review for so long.

The fix

hip_utils.resolve_target_archs() — explicit argument → FLASHINFER_ROCM_ARCH_LIST → architectures actually present → gfx942 with a warning. validate_rocm_arch, validate_flashinfer_rocm_arch, CompilationContext and aot_hip all route through it. _auto_detect_archs is deleted (private, no other caller, and the holder of the remaining literals).

The GPU-less fallback stays on gfx942 — the pre-existing default — so this changes nothing on hosts where no GPU is visible. aot_hip publishes the resolved list into FLASHINFER_ROCM_ARCH_LIST and resolve_aiter_build_arch() reads it back, so the two must agree about that default or the AITER shim is built for a different architecture than the kernels it ships beside. A test pins _GPULESS_FALLBACK_ARCH == aiter_source._DEFAULT_BUILD_ARCH so they cannot drift apart silently.

Detection uses rocminfo, not torch.cuda, so the resolver adds no torch dependency to a module that must stay importable without one — the hardware-less conformance job from #287 loads it.

Also included, both small and both fixing real failures:

• _canonical_arch_list — the validators split on , only and match tokens verbatim, so FLASHINFER_ROCM_ARCH_LIST=gfx942;gfx950 became one token and raised "does not support any of the requested ROCm architectures". ; is worth accepting because jit/aiter_source.py already documents it for this same variable — the two consumers were disagreeing about the format of their own env var. Qualifiers (gfx950:sramecc+), empty tokens and duplicates are handled too. Syntax only: unknown architectures pass through so the validators can still report them.
• Validate before publishing — aot_hip now validates the resolved list before writing it to the environment, so a failed build leaves the environment as it found it.

Test plan

• After the fix, all four resolution paths report gfx950 on the gfx950 host — the measurement at the top, re-run.
• FLASHINFER_ROCM_ARCH_LIST=gfx942 on a gfx950 box still resolves to gfx942; cross-compiling stays possible.
• A container with no /dev/kfd resolves to gfx942 with the warning, without importing torch.
• Detection is cached per process — CompilationContext is constructed from five places, and the code this replaces reached rocminfo through the cached get_supported_device_indices. Calling the uncached probe directly re-ran the subprocess, timeout included, on every construction. Covered by a test.
• test_defaults_to_gfx942_when_no_env_and_no_arg now pins the detected architecture instead of the literal.
• 195 passed across test_hip_utils.py, test_aot_hip.py, test_aiter_build_arch_hip.py, test_arch_caps_hip.py on gfx950.
• pre-commit run clean on all changed files.
• Full suite on gfx950 and gfx942 — to be run on merged amd-integration, same image both sides, so architecture is the only variable.

[Copilot]

Pull request overview

This PR centralizes ROCm target-architecture resolution into hip_utils.resolve_target_archs() so all build/validation paths (JIT validation, CompilationContext, and AOT packaging) agree on the same --offload-arch set—fixing a prior CDNA4 mismatch where validation could check gfx942 while compilation targeted gfx950.

Changes:

• Added resolve_target_archs() and routed validate_rocm_arch, validate_flashinfer_rocm_arch, CompilationContext, and aot_hip through it.
• Removed CompilationContext._auto_detect_archs and switched detection to the shared resolver.
• Updated/added tests to cover resolution branches and to assert validator/compile-context agreement.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
flashinfer/hip_utils.py	Introduces the unified architecture resolver and wires validators through it.
flashinfer/compilation_context_hip.py	Uses the unified resolver to avoid divergence between compilation and validation.
flashinfer/aot_hip.py	Resolves+publishes the target arch list once, then validates via CompilationContext.
tests/rocm_tests/test_hip_utils.py	Adds targeted tests for the new resolver and updates the prior “defaults to gfx942” assertion.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

flashinfer/aot_hip.py:231

• compile_and_package_modules() sets FLASHINFER_ROCM_ARCH_LIST before validating it via CompilationContext(). If CompilationContext() raises (e.g. unsupported ROCm version / arch), the process environment is left mutated to an invalid value, which can affect subsequent calls in the same process.

    rocm_arch_list = resolve_target_archs()
    os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list
    CompilationContext()  # validates the resolved list, raising on a bad one
    if verbose:
        print(f"Target ROCm architectures: {rocm_arch_list}")

[demandal25]

Suppressed comment in review 4984984345 — accepted, fixed in 4d1335d.

A CompilationContext() raise left FLASHINFER_ROCM_ARCH_LIST set for the rest of the process; now restored on failure. One correction: the leaked value was canonical, not invalid — a valid list that failed validation, which is why it was subtle.

My first regression test was vacuous (it seeded the variable with the value the resolver returns, so the buggy write was a no-op); rewritten and A/B-ed against the unfixed code.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (1)

flashinfer/aot_hip.py:234

• The comment says the arch list is “resolved once, then validated”, but CompilationContext() re-resolves via resolve_target_archs() (with the env var still unset), so this can re-run detection (extra rocminfo subprocess + duplicate warning) and the validation isn’t guaranteed to be against the exact rocm_arch_list you just computed. Consider temporarily setting FLASHINFER_ROCM_ARCH_LIST to rocm_arch_list just for the validation call, and restoring it on failure; on success you can leave it set to the resolved value.

    # the inputs cannot change in between, so it validates exactly the list
    # published below -- the order costs nothing and stops a raise here from
    # leaving FLASHINFER_ROCM_ARCH_LIST set for whatever runs next in-process.
    rocm_arch_list = resolve_target_archs()
    CompilationContext()  # validates the resolved list, raising on a bad one

[demandal25]

Suppressed comment in review 4985096141 — accepted, fixed in 6e8773c. This was a regression I introduced in 4d1335d, not pre-existing.

Both halves verified: rocminfo_gpu_agents() is not cached (only the two get_*_device_indices helpers are), and two resolves on a GPU-less host emit 2 warnings. Now publishes before validating and restores on failure, so detection runs once. The test pins both properties.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (2)

flashinfer/aot_hip.py:247

• compile_and_package_modules() publishes the resolved arch list to FLASHINFER_ROCM_ARCH_LIST, but CompilationContext() may silently filter that list (unsupported by ROCm version / unsupported by FlashInfer) without raising. If that happens, the env var can still contain unsupported entries (and even keep them first), which can cause the AITER shim (jit/aiter_source.py) to build for a different/invalid architecture than the kernels being packaged. Update the env var to the validated set after CompilationContext succeeds so the process-global side effect reflects what will actually be compiled.

    rocm_arch_list = resolve_target_archs()
    previous_arch_list = os.environ.get("FLASHINFER_ROCM_ARCH_LIST")
    os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list
    try:
        CompilationContext()  # validates the resolved list, raising on a bad one
    except BaseException:
        if previous_arch_list is None:
            os.environ.pop("FLASHINFER_ROCM_ARCH_LIST", None)
        else:
            os.environ["FLASHINFER_ROCM_ARCH_LIST"] = previous_arch_list
        raise

tests/rocm_tests/test_aot_hip.py:198

• This test creates a temporary build_dir via tempfile.mkdtemp() but never cleans it up, which can leak disk space across repeated test runs. Mirror test_compile_and_package_minimal() by deleting the directory in a finally block.

    with pytest.raises(RuntimeError, match="not recognized"):
        aot_hip.compile_and_package_modules(
            out_dir=None,
            build_dir=Path(tempfile.mkdtemp()),
            project_root=Path(__file__).parent.parent,
            config={
                "fa2_head_dim": [(128, 128)],
                "f16_dtype": [torch.float16],
                "use_sliding_window": [False],
                "use_logits_soft_cap": [False],
            },
            verbose=False,
            skip_prebuilt=True,
        )

[demandal25]

Suppressed comments in review 4985222187 — both accepted, fixed in 7bdc2c9. The first is the best catch on this PR.

• Validation filters rather than raises (gfx900,gfx942 → gfx942, warning only), so the variable could advertise an architecture the packaged kernels were never built for — and the AITER shim resolves its own target from it. Now republishes the validated list, ordered via arch_flags since TARGET_ROCM_ARCHS is a set.
• Leaked mkdtemp() → tmp_path fixture.

Both new tests were A/B-ed against the unfixed code.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Suppressed comments (4)

flashinfer/hip_utils.py:98

• This set comprehension includes every physical supported agent reported by rocminfo, even after the process has been restricted with HIP_VISIBLE_DEVICES. The repository documents that rocminfo ignores that variable (tests/rocm_tests/conftest.py:97-102); on a mixed gfx942/gfx950 host with HIP_VISIBLE_DEVICES selecting gfx950 and a gfx950-only PyTorch wheel, this resolves to both architectures and validate_flashinfer_rocm_arch then fails because PyTorch lacks gfx942. The resolver needs to honor the visible/current device when one is selected, or explicitly require an override for this case.

    detected = sorted(
        {
            arch
            for arch, _ in rocminfo_gpu_agents()
            if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS
        }
    )

flashinfer/hip_utils.py:107

• When FLASHINFER_ROCM_ARCH_LIST is set to a non-empty value that normalizes away (for example ";;" or whitespace), this warning says the variable is "unset" even though it was supplied. That misdirects the operator toward setting a variable that is already present; report that it is unset or contains no usable architecture instead.

    logger.warning(
        "No supported AMD GPU detected and FLASHINFER_ROCM_ARCH_LIST is unset; "
        "building for every supported architecture (%s). This is slower than "
        "targeting one. Set FLASHINFER_ROCM_ARCH_LIST to the architecture you "
        "are building for.",

tests/rocm_tests/test_hip_utils.py:345

• This new test covers validate_rocm_arch, but the existing TestValidateFlashinferRocmArch.test_defaults_to_gfx942_when_no_env_no_arg still mocks validate_rocm_arch and asserts the old hard-coded gfx942. A regression in validate_flashinfer_rocm_arch(arch_list=None) resolving the wrong target would therefore still pass. Update that wrapper test to stub rocminfo_gpu_agents and assert the detected target instead.

    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.

flashinfer/hip_utils.py:98

• When no environment override is present, this path starts a new rocminfo subprocess on every resolve_target_archs() call. The removed _auto_detect_archs() used the cached get_supported_device_indices(), but JIT setup constructs CompilationContext more than once and gen_jit_spec() revalidates for each operation, so a normal multi-op process now repeatedly pays this probe (including its 10-second timeout) and can emit the fallback warning repeatedly. Cache the hardware probe for the process or reuse the already-resolved list while preserving invalidation when visibility changes.

    detected = sorted(
        {
            arch
            for arch, _ in rocminfo_gpu_agents()
            if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS
        }
    )

[demandal25]

Review 4985367370 — 4 suppressed comments; 3 fixed in b2f586d, 1 deferred.

• rocminfo_gpu_agents() now cached. gen_jit_spec validates per module, so this was one subprocess (10 s timeout) per operation — a regression this PR introduced by dropping _auto_detect_archs(), which went through the already-cached get_supported_device_indices().
• Fallback warning no longer says "unset" when the variable was set but normalized away.
• test_defaults_to_gfx942_when_no_env_no_arg stubbed the resolved list out of its own assertion; it now echoes the argument and checks the detected arch.
• Deferred: rocminfo ignoring HIP_VISIBLE_DEVICES. Real, but honouring visibility is a behaviour change to detection, not to this PR's "resolve in one place". Same bucket as the two threads above.

164 passed; pre-commit clean.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (8)

flashinfer/aot_hip.py:235

• The no-environment behavior changed here, but amd-flashinfer-jit-cache/README.md:35-40 still says that AOT always compiles gfx942 and that the environment default is gfx942. A gfx950 host now follows detection, while a GPU-less host targets every supported architecture, so the package build instructions are inaccurate. Update that documentation (and the stale JIT default comment) with the centralized resolver behavior.

    rocm_arch_list = resolve_target_archs()

flashinfer/aot_hip.py:260

• context is used only to validate and derive the environment value; the actual HIP flags later come from the module-global jit.core.current_compilation_context (jit/core.py:406), which was initialized when flashinfer.jit was first imported. If this API is called after changing FLASHINFER_ROCM_ARCH_LIST (or called twice with different targets), context can validate/republish gfx950 while the kernels still compile with stale gfx942 flags; AITER then follows the republished value and no longer matches the packaged kernels. The active JIT context must be rebuilt/updated for this target before generating specs, or target changes must be rejected.

    rocm_arch_list = ",".join(
        flag.removeprefix("--offload-arch=") for flag in context.arch_flags
    )
    os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list

flashinfer/hip_utils.py:98

• The deleted _auto_detect_archs() caught all probe failures and fell back, but this new direct call is not protected. rocminfo_gpu_agents() only converts FileNotFoundError and timeouts to an empty result, so an unavailable executable that raises another OSError (for example PermissionError) now aborts resolution instead of reaching the documented GPU-less fallback. Catch the probe's OSError here or broaden the helper's failure handling.

    detected = sorted(
        {
            arch
            for arch, _ in rocminfo_gpu_agents()
            if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS
        }
    )

flashinfer/hip_utils.py:64

• The no-argument behavior documented by this new resolver is now detection/env dependent (and can fall back to both supported architectures), but flashinfer/jit/core.py:148 still says arch_list=None “defaults to gfx942”. That comment is on a call site now routed through this resolver, so it should be updated in this change to avoid documenting the behavior this PR removes.

    Step 4 replaces a hard-coded ``"gfx942"`` that three call sites reached
    independently. On a CDNA4 host that literal was not a conservative default
    but a wrong answer: ``validate_flashinfer_rocm_arch(arch_list=None)``
    returned ``{"gfx942"}`` on a gfx950 device while ``CompilationContext``
    compiled for gfx950, so the check that exists to catch "your PyTorch was not
    built for this architecture" was validating an architecture nobody was
    building for. Vacuous on a PyTorch carrying both; a spurious hard failure on
    an arch-specific build that carries only gfx950.

tests/rocm_tests/test_aot_hip.py:213

• This new test hard-codes gfx950 but runs CompilationContext against the real PyTorch extension flags. On a supported gfx942 host with an arch-specific PyTorch build, validation filters gfx900 and then raises because PyTorch does not advertise gfx950, so the test fails before checking the AOT environment behavior. Stub _get_rocm_arch_flags (as the resolver-agreement test does) or choose the host's target so this test is hermetic.

    monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx900,gfx950")

flashinfer/aot_hip.py:231

• The explanation here is now stale: rocminfo_gpu_agents() is decorated with functools.cache, and publishing the environment before constructing CompilationContext means its resolver takes the just-published value rather than performing a second detection pass. Please update the comment so it does not claim an uncached probe or a second warning.

    # Publish first so CompilationContext() -- which resolves through
    # resolve_target_archs() itself -- validates exactly this list rather than
    # repeating the work. rocminfo_gpu_agents() is not cached, so simply
    # validating before publishing would re-run detection and, on a GPU-less
    # host, emit the "no supported AMD GPU detected" warning a second time.

flashinfer/hip_utils.py:505

• The new @functools.cache makes this probe process-wide, but the docstring below still says Not cached: the caller decides and describes callers paying for a fresh subprocess. That is now false and can lead callers to reason incorrectly about cache invalidation; update the stale sentence to document the shared cache and how it is cleared.

    Cached for the process. ``gen_jit_spec`` calls ``check_rocm_arch()`` for
    every operation it builds, and that reaches here through
    ``resolve_target_archs()``, so without this a multi-op process pays one
    ``rocminfo`` subprocess -- and its 10-second timeout -- per module, and a
    GPU-less host repeats the fallback warning each time. The prior

tests/rocm_tests/test_aot_hip.py:196

• This rationale is stale after rocminfo_gpu_agents became cached in the same change: CompilationContext will not perform a second rocminfo subprocess merely because it resolves again. Keep the assertion, but describe the actual invariant—that validation must receive the exact resolved list rather than re-reading an unset environment variable—so the test does not document behavior that no longer exists.

    # Validation must see the resolved list, not an unset variable. Otherwise
    # CompilationContext re-resolves from scratch -- rocminfo_gpu_agents() is
    # not cached, so that is a second detection pass and, on a GPU-less host, a
    # second "no supported AMD GPU detected" warning.

[demandal25]

Review 4985855053 — 8 suppressed; 7 fixed in b3fd023, 1 deferred.

Four were defects the previous commit introduced. Adding @functools.cache to rocminfo_gpu_agents falsified three of my own comments (the leftover "Not cached: the caller decides" paragraph, and the "avoids a second detection pass" rationale in both aot_hip and its test) — all rewritten to the reason that survives caching. The fourth: my new test_environment_gets_the_validated_list_not_the_resolved_one repeated the exact hermeticity bug flagged earlier on a different test, hard-coding gfx950 against the real PyTorch. It now stubs _get_rocm_arch_flags; verified with a plugin simulating a gfx942-only wheel that it and the other two arch-sensitive tests pass.

Three were pre-existing and surfaced because this PR changes what they describe: jit/core.py's "defaults to gfx942" comment, amd-flashinfer-jit-cache/README.md, and rocminfo_gpu_agents catching only FileNotFoundError/TimeoutExpired so an unexecutable rocminfo (PermissionError) aborted resolution instead of reaching the GPU-less fallback.

Deferred: the module-global JIT context (aot_hip.py:260) — same finding as the thread above, unchanged in scope.

164 passed; pre-commit clean.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (4)

flashinfer/hip_utils.py:341

• validate_rocm_arch only consults resolve_target_archs when arch_list is None, so an explicit caller value bypasses the new canonicalization. For example, validate_rocm_arch("gfx950:sramecc+") reaches the compatibility matrix with the qualifier and fails even though resolve_target_archs accepts and normalizes caller-supplied lists. Resolve unconditionally here so explicit and implicit inputs use the same path.

    if arch_list is None:
        arch_list = resolve_target_archs()

flashinfer/hip_utils.py:425

• The wrapper has the same bypass: when a caller supplies arch_list, it passes the raw value straight to validate_rocm_arch instead of using the centralized resolver. Thus validate_flashinfer_rocm_arch("gfx942;gfx950") still treats the semicolon list as one invalid token, while the resolver's documented canonical path handles it. Resolve the argument unconditionally before validation.

    # Get architecture list from parameter, env var, or default
    if arch_list is None:
        arch_list = resolve_target_archs()

tests/rocm_tests/test_aot_hip.py:214

• This test depends on the host's real ROCm version. On the supported ROCm 6.4/6.3 path (_ROCM_ARCH_GROUPS above), neither gfx900 nor gfx950 is supported, so CompilationContext() raises before the republished environment is checked instead of filtering to gfx950. Stub get_system_rocm_version to a 7.x version here, as the architecture tests do, so the test is portable across the supported ROCm matrix.

    monkeypatch.setenv("FLASHINFER_ROCM_ARCH_LIST", "gfx900,gfx950")

flashinfer/hip_utils.py:117

• On a GPU-less host with no environment override, this warning is emitted on every resolve_target_archs() call; caching rocminfo_gpu_agents() does not cache this fallback branch. gen_jit_spec() calls check_rocm_arch() for each operation, so one JIT generation can flood stderr/logs with the same warning. Guard the fallback warning (or cache the no-device resolution while still invalidating it when the environment changes) so the operator is notified once.

    logger.warning(
        "No supported AMD GPU detected and %s; building for every supported "
        "architecture (%s). This is slower than targeting one. Set "
        "FLASHINFER_ROCM_ARCH_LIST to the architecture you are building for.",
        how,
        fallback,
    )

[demandal25]

Suppressed comments from review 4996642977 — all four read, three fixed in c5e5bea.

• hip_utils.py:83 (UUID selectors) — fixed. rocminfo does report a Uuid: per agent, so GPU-<uuid> now maps to the same enumeration as an index. Reviewing this fix also turned up a bug it introduced: the empty-string placeholder stored for a uuid-less agent matched HIP_VISIBLE_DEVICES="", reporting a GPU as visible for the canonical hide-everything form. Reproduced and fixed too.
• hip_utils.py:150 (empty env var reported as "is not set") — fixed. I had changed is None to a falsiness test; restored.
• aiter_source.py:77 (env leak on a raise) — reordered, but the leak is not reachable. _aiter_libs_dir() resolves before any write, so a raise always precedes AITER_SYMBOL_VISIBLE/AITER_JIT_DIR. I moved the resolve above the writes anyway rather than rely on that. No test: the only test I could write for it passed with and without the fix, so I dropped it rather than keep a vacuous one.
• jit/env.py:216 (description vs. the new import behaviour) — description was stale, updated ~1 min before this review landed. It now covers the GPU-less import change and why _get_workspace_dir_name needed the narrower exception type.

Separately: I had pushed the redesign without running the quality gate's code review. Running it found eight defects, four of mine, two reproduced on-box — all fixed in c5e5bea.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

flashinfer/jit/core.py:183

• Because this is a process-global lazy singleton, two concurrent first accesses can both observe _ctx is None while CompilationContext() is doing its probe and validation. That runs the supposedly one-time resolution more than once (and the last assignment wins if the device/environment changes), so the proxy is not actually a once-per-process context. Guard the check-and-construct with a lock or another thread-safe lazy-initialization primitive.

        if self._ctx is None:
            self._ctx = CompilationContext()

[demandal25]

Suppressed comment from review 4997466925 — accepted, fixed in 19fc3e5.

jit/core.py:183 (lazy singleton race) — correct, and it is a regression this branch introduced: constructing the context at import meant Python's import lock serialized it, and deferring moved it out from behind that lock. Two threads reaching a first JIT build together would each build a context, so "resolved once per process" did not hold. Double-checked locking on self._lock; the fast path stays lock-free.

Separately, I had pushed c5e5bea without the self code-review the quality gate requires. Running it on 6e58e0d..HEAD found five more, two serious — both mine, both fixed in cd390d5:

• The workspace key I added in 19fc3e5 built a directory name from FLASHINFER_ROCM_ARCH_LIST with no validation. =/tmp/pwn made FLASHINFER_WORKSPACE_DIR exactly /tmp/pwn (pathlib discards everything left of an absolute component), and jit/core.py mkdirs that at import. Same traversal class this PR already closed in _aiter_cache_tag.
• tests/conftest.py contradicted its own comment: the empty-probe branch pinned the raw worker index, which is the out-of-range case the comment two lines above says is broken.

Plus a docstring that overclaimed, dead _ARCH_RE, and two weak tests (a pytest.warns with no match=, and one that passed with the fix deleted).

259 tests pass serially and at -n 2 on a rebuilt rocm7.2 image.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

flashinfer/jit/aiter_source.py:88

• The cache key omits the value that this function uses to choose the result: _detected_device_arch() at lines 97-99. With both gfx942 and gfx950 in the resolved list, the first call on a gfx942 device caches gfx942; after torch.cuda.set_device() switches to gfx950, the resolved string is unchanged and this cache still selects gfx942, so the AITER shim/link path can be reused on the wrong GPU. Include the current device architecture in the selection key or do not cache this device-dependent choice.

@functools.lru_cache(maxsize=None)
def _select_build_arch(resolved: 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 the resolved string rather than ``maxsize=1``, so a changed target
    re-selects, and each distinct list warns once rather than once per module.

flashinfer/hip_utils.py:161

• visible_gpu_agents() has already established that the process has no visible device, but this unconditional torch fallback can resurrect one if the HIP runtime was initialized before HIP_VISIBLE_DEVICES/CUDA_VISIBLE_DEVICES was changed. For example, with HIP_VISIBLE_DEVICES=-1 (or an empty/invalid selector), a previously initialized torch.cuda can still report its devices, causing the resolver to return a target instead of raising. Make the torch fallback honor explicit visibility (at least the hide-all/invalid forms), and cover that path without stubbing _torch_detected_archs.

    detected = _torch_detected_archs()
    if detected:
        return ",".join(detected)

flashinfer/jit/core.py:196

• This lazy proxy is now resolved after jit/env.py has already fixed FLASHINFER_WORKSPACE_DIR from the import-time current device. On a mixed host, tests/conftest.py can pin the worker to gfx950 after that import, so the first access here emits gfx950 flags into a workspace keyed as gfx942; the same mismatch occurs if the target/device changes after import. That lets a cache directory reuse or publish binaries under the wrong architecture. The workspace key must be deferred/refreshed together with the realized context (or this deferral must wait until workspace selection is aligned).

current_compilation_context = _LazyCompilationContext()

[demandal25]

Suppressed comments from review 4997770801 — all three read; one fixed, two answered.

• aiter_source.py:88 (cache key omits the device) — fixed in 21e7719, and fixing it needed a second pass. Adding the device to the key made the answer track an uncached probe, so it could differ between the two calls one AITER module makes: ensure_aiter_lib copies the .so into the directory the first names, aiter_jitspec_flags builds -L/-rpath from the second. _detected_device_arch is cached now, so the pair is stable — an unstable answer is worse than a frozen one. Regression test added.
• hip_utils.py:161 (torch fallback vs explicit visibility) — tried and reverted, see 24957ed. Gating on "rocminfo itself found nothing" skipped the fallback whenever visibility parsing failed (a GPU-<uuid> selector against a rocminfo that omits Uuid — a working host), and re-ran a probe that is deliberately not cached when empty, doubling a 10 s timeout. torch.cuda.device_count() honours HIP_VISIBLE_DEVICES itself, which is what makes the ungated call safe; the residual case needs a runtime initialized before the variable changed, and there torch reporting the machine's real architecture is not a wrong answer. Happy to revisit if you have a case that survives that.
• jit/core.py:196 (workspace fixed at import vs lazy target) — correct, not fixed. Same root as the two _get_workspace_dir_name threads; the real fix is deferring the workspace alongside the context, which is a wide change across module-level Paths that aot_hip reassigns. Tracked as the follow-up — partial patches in this area have generated more defects than they fixed.

Process note: this round's fixes were themselves self-reviewed before pushing, which is how the two corrections above were caught rather than shipped.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

flashinfer/jit/aiter_source.py:257

• This resolves the target a second time after _aiter_libs_dir() has already derived the cache tag and lib_path. The outer resolver is no longer cached and rereads FLASHINFER_ROCM_ARCH_LIST, so if the target changes between these calls, the path can be for architecture A while GPU_ARCHS builds architecture B; B is then copied into A's cache directory and linked through A's RUNPATH. Resolve once and pass the same architecture/tag through both path construction and the AITER build environment.

    gpu_archs = resolve_aiter_build_arch()

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated no new comments.

Suppressed comments (3)

flashinfer/jit/aiter_source.py:90

• On a host where rocminfo is unavailable or times out but the documented _torch_detected_archs() fallback succeeds, this call still retries the 10-second rocminfo probe on every AITER module: rocminfo_gpu_agents() deliberately does not cache empty results, and resolve_target_archs() runs it before the torch fallback. aiter_jitspec_flags() is invoked independently by the activation, norm, and rope AITER generators, so one JIT setup can pay this delay several times. Cache/share the resolved target (or otherwise cache a failed probe separately from the initial worker-discovery path) before resolving per module.

    return _select_build_arch(resolve_target_archs(), _detected_device_arch())

flashinfer/jit/aiter_source.py:107

• This device-aware cache is bypassed by the normal AITER entry points: get_silu_and_mul_aiter_module, get_rope_aiter_module, get_norm_aiter_module, and get_batch_decode_aiter_module are all @functools.cached only by their shape/dtype arguments (for example, activation.py:38 and decode_rocm.py:300). After the first module is built on gfx942, switching the current device to gfx950 returns the already-loaded shim without calling resolve_aiter_build_arch() again, so the old architecture-specific library/RUNPATH is still used. Include the device/architecture in those public cache keys or invalidate those module caches when the device changes.

    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.

flashinfer/jit/env.py:292

• _arch_cache_key() is only evaluated while jit.env is imported, before _LazyCompilationContext resolves the target. On a GPU-less import with no env this therefore fixes FLASHINFER_WORKSPACE_DIR to .../noarch; setting FLASHINFER_ROCM_ARCH_LIST=gfx950 before the first JIT build still compiles into that directory, and a later gfx942 build can reuse the same build.ninja/.so. Defer workspace selection until the target is realized, or re-key/invalidate the workspace when the context resolves.

            arch = _arch_cache_key(os.environ.get("FLASHINFER_ROCM_ARCH_LIST", ""))

[demandal25]

Suppressed comments from review 4998975841.

jit/env.py:292 (workspace vs lazy target) — fixed, 5a2f659 + e372913. Measured before: FLASHINFER_ROCM_ARCH_LIST=gfx942 compiled gfx942 into .../gfx950/cached_ops. An explicit target now names the directory. Caveat: it is the requested list, which validation may narrow, so it is an improvement rather than an invariant.

aiter_source.py:107 (module caches keyed on shape/dtype) — declined. A loaded .so cannot be unloaded, so an already-loaded module is the only thing those caches can return; a device-keyed entry would build a second shim that cannot replace the first.

aiter_source.py:90 (repeated 10 s probe) — real, left unfixed. I memoised the resolver and reverted it: functools.cache does not cache exceptions, so the expensive path was unhelped, and the key was the environment while the answer comes from probes every test stubs — a stubbed arch leaked into production resolution. Wrong layer; the cost belongs at the probe.

[Copilot]

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

flashinfer/jit/aiter_source.py:114

• This candidate list is validated only against ROCm and FlashInfer because no PyTorch extension module is supplied. In the normal AITER path, aiter_jitspec_flags() runs before gen_jit_spec(), so an explicit/multi-architecture target missing from the installed PyTorch can spend time building an AITER shim and only then fail when the lazy compilation context performs the PyTorch check. Select from the shared context's already-validated flags, or perform the PyTorch validation before ensure_aiter_lib().

    arch_flags, _ = validate_flashinfer_rocm_arch(arch_list=resolved)
    candidates = [f.removeprefix("--offload-arch=") for f in arch_flags]

flashinfer/aot_hip.py:198

• This validates with a new CompilationContext, but the generated specs do not use this object: gen_jit_spec() later realizes current_compilation_context and reads its flags. AOT therefore resolves/validates twice, and the two objects can disagree if visibility, device, or the environment changes between these calls. Realize the shared proxy here so the early check and the context that emits --offload-arch are the same object.

    from .compilation_context_hip import CompilationContext

    context = CompilationContext()  # raises if nothing in the list is usable

flashinfer/jit/aiter_source.py:90

• _select_build_arch is cached, but its arguments are evaluated first, so every resolve_aiter_build_arch() call still invokes resolve_target_archs(). When rocminfo is unavailable or times out, rocminfo_gpu_agents() deliberately does not cache the empty result; consequently each AITER module retries a 10-second probe even though _torch_detected_archs() already supplied the target. On the documented slim-image/torch-fallback path, this can add tens of seconds to one import/build. Cache the resolved probe/result for the operation (while preserving the intentional retry semantics for the initial visibility setup), or otherwise avoid re-probing for each AITER module.

    return _select_build_arch(resolve_target_archs(), _detected_device_arch())

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (6)

flashinfer/aot_hip.py:234

• CompilationContext() can validate a narrowed target list, but this publishes the pre-validation string. For example, on ROCm 6.4 gfx950,gfx942 is narrowed to gfx942; resolve_aiter_build_arch() then sees the published fat list and can select gfx950 on a gfx950 device while the FlashInfer kernels target only gfx942. Publish the validated flags/targets from the context so the AITER shim receives the same effective architecture set.

    rocm_arch_list = resolve_target_archs()
    CompilationContext()  # validates the resolved list, raising on a bad one
    os.environ["FLASHINFER_ROCM_ARCH_LIST"] = rocm_arch_list

flashinfer/aot_hip.py:233

• CompilationContext() is invoked only after FLASHINFER_WORKSPACE_BASE and the jit_env workspace paths have been mutated above. If this validation raises, the process retains the failed AOT build_dir as its global JIT workspace, so subsequent builds or retries use the wrong directory. Move resolution/validation before those global overrides and before importing the JIT modules.

    CompilationContext()  # validates the resolved list, raising on a bad one

flashinfer/hip_utils.py:45

• rocminfo_gpu_agents() enumerates every physical GPU and does not apply HIP_VISIBLE_DEVICES. Because this new cache is used by CompilationContext during import, a worker pinned to gfx950 on a mixed gfx942/gfx950 host still resolves both targets; an arch-specific PyTorch build can then fail validation, while a fat build compiles code for an invisible GPU. Filter the probe by the process-visible device selection before caching, or defer resolution until the visibility pin is established.

                for arch, _ in rocminfo_gpu_agents()
                if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS

flashinfer/hip_utils.py:143

• When the environment is present but normalizes entirely away (";;" or whitespace)—a case this resolver deliberately supports—this branch is taken even though the warning says FLASHINFER_ROCM_ARCH_LIST is unset. The message should say that no usable target was provided (or identify the invalid value), otherwise it misdiagnoses the operator's configuration.

    logger.warning(
        "No supported AMD GPU detected and FLASHINFER_ROCM_ARCH_LIST is unset; "
        "falling back to %s. Set FLASHINFER_ROCM_ARCH_LIST to the architecture "
        "you are building for -- otherwise the result will not run on %s.",

flashinfer/hip_utils.py:44

• The removed _auto_detect_archs() caught probe failures and used the fallback, but this new wrapper calls rocminfo_gpu_agents() without protection. That helper only handles FileNotFoundError and TimeoutExpired, so a PermissionError/other OSError from subprocess.run now aborts CompilationContext (and potentially package import) instead of taking the documented fallback path. Catch probe failures here and return an empty tuple.

                for arch, _ in rocminfo_gpu_agents()

tests/rocm_tests/test_hip_utils.py:384

• _detected_supported_archs is process-cached, but the cache-clearing fixture is scoped only to TestResolveTargetArchs. If the resolver ran during package import or an earlier test, patching rocminfo_gpu_agents here is ignored; running this test alone (or after another module) on a gfx942 host can still return gfx942 and fail the new gfx950 assertion. Clear this cache before and after the patched call, or use a module-scoped fixture.

            patch(
                "flashinfer.hip_utils.rocminfo_gpu_agents",
                return_value=(("gfx950", "AMD Instinct MI350X"),),

[demandal25]

The six Copilot comments above predate the scope reduction (they were written 2026-08-20/21; the branch was rewritten on 08-22) and none apply to the current diff. Verified rather than assumed:

comment	status
visible_gpu_agents() scoping	symbol no longer exists
resolve_target_archs() raising TargetArchUnresolved	no longer raises; falls back to gfx942, the pre-existing default
_torch_detected_archs() observing an unpinned runtime	symbol no longer exists
aiter_source.py _ARCH_RE candidate selection	file is not in this PR
jit/core.py lazy-context init ordering	file is not in this PR
jit/env.py noarch workspace sharing one JIT cache	file is not in this PR

One is now inverted by the rewrite: the aot_hip comment says the hunk removes the publication into FLASHINFER_ROCM_ARCH_LIST. It does not — aot_hip.py:234 still publishes it, and the PR description describes that as the load-bearing AOT → AITER channel.

The current diff is 5 files / +414 −57, confined to hip_utils.py, compilation_context_hip.py, aot_hip.py and their tests. 195 tests pass; both checks are green.

To be clear about what is being set aside rather than dismissed: the jit/env.py cache-key observation and the "HIP_VISIBLE_DEVICES is set after the package import" observation describe real defects. Both are pre-existing and independent of this change — the latter is a property of jit/env.py querying device properties at module scope, which this PR does not touch. Fixing either means restructuring import-time initialization, which is a larger change than the one-line-default bug in the title, and folding them in is what took this PR from 1 commit to 32 without converging.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Suppressed comments (3)

flashinfer/hip_utils.py:143

• When the variable is set to a non-empty value that normalizes away (for example " ; , ", which this PR explicitly supports as a fallback case), this branch still logs that FLASHINFER_ROCM_ARCH_LIST is unset. That makes the warning misleading and hides the invalid/empty configuration; distinguish unset from an empty-after-normalization value or describe both cases.

        "No supported AMD GPU detected and FLASHINFER_ROCM_ARCH_LIST is unset; "
        "falling back to %s. Set FLASHINFER_ROCM_ARCH_LIST to the architecture "
        "you are building for -- otherwise the result will not run on %s.",

flashinfer/hip_utils.py:45

• This cache is independent of the existing get_supported_device_indices() cache, so it does not prevent duplicate probes. tests/conftest.py calls get_physical_card_device_indices() (which invokes rocminfo_gpu_agents()), and the new resolver invokes the uncached probe again in the same worker; with the 10-second subprocess timeout this doubles startup/probe cost. Share a cached raw-agent result between these helpers instead of caching only the derived architecture set here.

    return tuple(
        sorted(
            {
                arch
                for arch, _ in rocminfo_gpu_agents()
                if arch in FLASHINFER_SUPPORTED_ROCM_ARCHS

tests/rocm_tests/test_aot_hip.py:174

• This test is vacuous against the parent implementation: when the variable is unset, the old compile_and_package_modules branch also constructed CompilationContext() before assigning FLASHINFER_ROCM_ARCH_LIST, so patching that constructor to raise already left the environment untouched. Replace this with an assertion that distinguishes the new resolver/publication behavior (and covers the validated target), otherwise a regression in the claimed ordering can pass this test.

    monkeypatch.setattr("flashinfer.compilation_context_hip.CompilationContext", _Boom)

    with pytest.raises(RuntimeError, match="not recognized"):
        aot_hip.compile_and_package_modules(
            out_dir=None,