fix(rocm): resolve the build architecture in one place, not three - #289
fix(rocm): resolve the build architecture in one place, not three#289demandal25 wants to merge 18 commits into
Conversation
Three call sites independently answered "what architecture are we building
for", and on CDNA4 they disagreed. Measured on a gfx950 host with
FLASHINFER_ROCM_ARCH_LIST unset, before this change:
ACTUAL device arch gfx950
validate_flashinfer_rocm_arch(arch_list=None) ['gfx942'] <-- wrong
CompilationContext().TARGET_ROCM_ARCHS ['gfx950']
resolve_aiter_build_arch() gfx950
So the JIT validated gfx942 while compiling for gfx950. The check exists to
catch "your PyTorch was not built for this architecture", and it was asking
about an architecture nobody was building for: vacuous on a PyTorch carrying
both, and a spurious hard failure on an arch-specific build carrying only
gfx950. The cause was `os.environ.get("FLASHINFER_ROCM_ARCH_LIST", "gfx942")`
reached from two functions in hip_utils, plus two more `return "gfx942"` lines
in CompilationContext._auto_detect_archs.
Add hip_utils.resolve_target_archs() -- explicit argument, then the env var,
then the architectures actually present, then every supported architecture with
a warning -- and route validate_rocm_arch, validate_flashinfer_rocm_arch,
CompilationContext and aot_hip through it. _auto_detect_archs goes away; it was
private and had no other caller.
The last-resort fallback changes from "gfx942" to every supported architecture.
On a GPU-less build host the old literal was not conservative, it was a guess
that silently produced a gfx942-only artifact; a fat build is slower but
correct wherever it lands, and the warning says how to make it cheap again.
Detection uses rocminfo rather than torch.cuda, so the resolver adds no torch
dependency to a module that must stay importable without one.
test_defaults_to_gfx942_when_no_env_and_no_arg asserted the wrong constant --
which is how the constant survived -- and now pins the detected architecture.
Measured after, same host and script: all four rows report gfx950. Also
verified FLASHINFER_ROCM_ARCH_LIST=gfx942 is still honoured on a gfx950 box
(cross-compiling stays possible), and that a container with no /dev/kfd
resolves to "gfx942,gfx950" with the warning and without importing torch.
169 passed across test_hip_utils.py, test_aiter_build_arch_hip.py and
test_arch_caps_hip.py on gfx950.
Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
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 routedvalidate_rocm_arch,validate_flashinfer_rocm_arch,CompilationContext, andaot_hipthrough it. - Removed
CompilationContext._auto_detect_archsand 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.
…t test hermetic Addresses both review comments on #289. - resolve_target_archs() returned the caller/env string verbatim. Now that it is the single source of truth that is a hard failure, not untidiness: the validators split on "," only and match tokens against FLASHINFER_SUPPORTED_ROCM_ARCHS verbatim, so FLASHINFER_ROCM_ARCH_LIST=gfx950:sramecc+ -> ['gfx950:sramecc+'] unsupported FLASHINFER_ROCM_ARCH_LIST=gfx942;gfx950 -> ['gfx942;gfx950'] unsupported FLASHINFER_ROCM_ARCH_LIST=gfx942,,gfx942 -> ['gfx942','','gfx942'] unsupported '' and validate_flashinfer_rocm_arch raises "does not support any of the requested ROCm architectures". ';' matters specifically because jit/aiter_source.py already documents it for this same variable, and aot_hip.py writes this resolver's output back into that env var -- so the two consumers were disagreeing about their own input format. _canonical_arch_list normalizes syntax only: accepts ',' or ';', strips qualifiers via normalize_arch, drops empties, dedupes preserving first-seen order. Unknown architectures pass through so the validators can still report them; dropping one here would turn a clear error into a build that quietly targets less than was asked for. A value that normalizes away entirely (";;") falls through to detection rather than returning "". - test_agrees_with_the_compilation_context compared a _FakeCppExt-fed validator against a CompilationContext that validates against the *real* torch, so the assertion depended on the installed wheel. Both sides now see the same view. Verified rather than assumed, by simulating an arch-specific wheel (_get_rocm_arch_flags -> gfx942 only) via a pytest plugin: before: RuntimeError: PyTorch does not support the following architectures: --offload-arch=gfx950 -> FAILED after: passes The wheel on this box is a fat build advertising gfx950, which is why the fragility was latent here. 11 tests added for the canonicalization. gfx942 behaviour is unchanged.
There was a problem hiding this comment.
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}")
…ment Addresses the suppressed comment in Copilot review 4984984345 on #289. compile_and_package_modules set FLASHINFER_ROCM_ARCH_LIST and then validated it, so a CompilationContext() raise left the variable set for whatever ran next in the same process. Publishing it is deliberate -- the AITER shim reads the build architecture from there (jit/aiter_source.py) and an AOT build has no other channel -- but it is a side effect that outlives the call, so it should only happen once the list is known good. Reordering is behaviour-preserving on success: CompilationContext() re-resolves through resolve_target_archs() and nothing can change between the two calls, so it validates exactly the list published afterwards. On the failure path the environment is now left as it was found. One correction to the review's wording: the leaked value is not "invalid" -- resolve_target_archs() returns a canonical list. It is a *valid* list that failed validation against this ROCm version or PyTorch build, which is a different thing and is why the leak is subtle rather than obvious. The regression test was wrong on its first attempt and is worth flagging: it seeded FLASHINFER_ROCM_ARCH_LIST with the value the resolver would return, so the buggy write was a no-op and the test passed with the bug present. It now starts from the variable unset with detection patched, and was verified to fail without the reorder: AssertionError: assert 'FLASHINFER_ROCM_ARCH_LIST' not in environ({... 'FLASHINFER_ROCM_ARCH_LIST': 'gfx950'})
|
Suppressed comment in review 4984984345 — accepted, fixed in 4d1335d. A 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. |
There was a problem hiding this comment.
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 viaresolve_target_archs()(with the env var still unset), so this can re-run detection (extrarocminfosubprocess + duplicate warning) and the validation isn’t guaranteed to be against the exactrocm_arch_listyou just computed. Consider temporarily settingFLASHINFER_ROCM_ARCH_LISTtorocm_arch_listjust 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
… failure Addresses the suppressed comment in Copilot review 4985096141 on #289, which caught a regression introduced by the previous commit. Validating before publishing made CompilationContext() re-resolve from an unset environment. rocminfo_gpu_agents() is not cached -- only get_supported_device_indices and get_physical_card_device_indices are -- so that is a second detection pass, and on a GPU-less host a second "no supported AMD GPU detected" warning. Measured: two resolves emit two warnings. Publishing first and restoring in an except clause gets both properties at once: validation sees exactly the list that was resolved (one detection, one warning), and a raise leaves the variable as it was found, including the case where it was previously unset. The test now pins both halves, since the leak fix alone passed without the single-resolution property: _Boom records what os.environ held when validation ran, so a return to validate-then-publish fails on `_Boom.seen == "gfx950"` rather than quietly reintroducing the duplicate detection.
|
Suppressed comment in review 4985096141 — accepted, fixed in 6e8773c. This was a regression I introduced in 4d1335d, not pre-existing. Both halves verified: |
There was a problem hiding this comment.
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,
)
Addresses both suppressed comments in Copilot review 4985222187 on #289. - Validation filters rather than raises. An architecture FlashInfer cannot serve is dropped with a warnings.warn and the build continues, provided at least one survives. Measured: validate_flashinfer_rocm_arch("gfx900,gfx942") -> no exception, arch_flags ['--offload-arch=gfx942'], set {'gfx942'} so FLASHINFER_ROCM_ARCH_LIST could advertise gfx900 while the packaged kernels were compiled only for gfx942. The AITER shim resolves its own build target from that variable (jit/aiter_source.py), so it would build for an architecture nothing else in the package targets -- the precise divergence this PR exists to remove, reintroduced one layer up. The resolved list is still published before validation, so CompilationContext does not re-run detection; the validated list replaces it afterwards. Taken from arch_flags rather than TARGET_ROCM_ARCHS because the latter is a set and order is meaningful, both on the hipcc command line and to AITER. - The new failure-path test leaked a tempfile.mkdtemp() directory. Switched to the tmp_path fixture, which pytest cleans up, rather than adding a finally block. Verified the guard is real, not vacuous -- without the republish: - gfx950 + gfx900,gfx950
|
Suppressed comments in review 4985222187 — both accepted, fixed in 7bdc2c9. The first is the best catch on this PR.
Both new tests were A/B-ed against the unfixed code. |
There was a problem hiding this comment.
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 withHIP_VISIBLE_DEVICES. The repository documents thatrocminfoignores that variable (tests/rocm_tests/conftest.py:97-102); on a mixed gfx942/gfx950 host withHIP_VISIBLE_DEVICESselecting gfx950 and a gfx950-only PyTorch wheel, this resolves to both architectures andvalidate_flashinfer_rocm_archthen 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_LISTis 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 existingTestValidateFlashinferRocmArch.test_defaults_to_gfx942_when_no_env_no_argstill mocksvalidate_rocm_archand asserts the old hard-codedgfx942. A regression invalidate_flashinfer_rocm_arch(arch_list=None)resolving the wrong target would therefore still pass. Update that wrapper test to stubrocminfo_gpu_agentsand 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
rocminfosubprocess on everyresolve_target_archs()call. The removed_auto_detect_archs()used the cachedget_supported_device_indices(), but JIT setup constructsCompilationContextmore than once andgen_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
}
)
…d a stale test Addresses three of the six findings in Copilot review 4985367370 on #289. The other three are design-level and tracked separately. - rocminfo_gpu_agents() is now cached. gen_jit_spec() calls check_rocm_arch() for every module it builds, and that reaches the probe through resolve_target_archs(), so a multi-op process was paying one rocminfo subprocess -- with a 10 s timeout -- per operation, and a GPU-less host repeated the fallback warning each time. The _auto_detect_archs() this PR removed avoided that by going through the already-cached get_supported_device_indices(), so this is a regression the PR introduced. Caching the probe itself also makes that function's "rocminfo is invoked at most once per process" docstring true globally rather than only on its path. TestGetSupportedDeviceIndices already cleared the derived cache per test; it now clears this one too. Without that the first test in the class pinned the probe and the remaining five asserted against it instead of their own patched subprocess output -- which is how they failed when the cache was added. - The GPU-less fallback warning said FLASHINFER_ROCM_ARCH_LIST "is unset" even when it was set to something that normalized away (";;", whitespace), which is reachable through the fall-through added earlier in this branch. It now distinguishes unset from "names no architecture", so the operator is not sent looking for a variable that is already there. - test_defaults_to_gfx942_when_no_env_no_arg stubbed validate_rocm_arch to return "gfx942" whatever it was handed, so the resolved list never reached the assertion: the hard-coded default this PR exists to remove could have come back with the test still green. The stub now echoes its argument and the detected architecture is what is checked.
|
Review 4985367370 — 4 suppressed comments; 3 fixed in b2f586d, 1 deferred.
164 passed; |
There was a problem hiding this comment.
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-40still 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
contextis used only to validate and derive the environment value; the actual HIP flags later come from the module-globaljit.core.current_compilation_context(jit/core.py:406), which was initialized whenflashinfer.jitwas first imported. If this API is called after changingFLASHINFER_ROCM_ARCH_LIST(or called twice with different targets),contextcan validate/republishgfx950while the kernels still compile with stalegfx942flags; 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 convertsFileNotFoundErrorand timeouts to an empty result, so an unavailable executable that raises anotherOSError(for examplePermissionError) now aborts resolution instead of reaching the documented GPU-less fallback. Catch the probe'sOSErrorhere 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:148still saysarch_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
gfx950but runsCompilationContextagainst the real PyTorch extension flags. On a supported gfx942 host with an arch-specific PyTorch build, validation filtersgfx900and then raises because PyTorch does not advertisegfx950, 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 withfunctools.cache, and publishing the environment before constructingCompilationContextmeans 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.cachemakes this probe process-wide, but the docstring below still saysNot cached: the caller decidesand 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_agentsbecame cached in the same change:CompilationContextwill 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.
… failure handling Addresses seven of the eight suppressed comments in Copilot review 4985855053 on #289. Four of them are defects the previous commit introduced. Self-inflicted by adding @functools.cache to rocminfo_gpu_agents: - The docstring still carried "Not cached: the caller decides" three paragraphs below the new "Cached for the process". Removed. - aot_hip's rationale for publishing before validating said it avoided "a second detection pass". True when written, false once the probe was cached. Rewritten to the reason that survives: the two agree by construction rather than by coincidence. - The same stale reasoning in the test comment, likewise rewritten. Also self-inflicted, and the same bug flagged earlier on a different test: - test_environment_gets_the_validated_list_not_the_resolved_one hard-coded gfx950 and validated against the real PyTorch, so on an arch-specific gfx942 wheel it would fail for reasons unrelated to the republish. Stubs _get_rocm_arch_flags now. Verified with a plugin simulating a gfx942-only wheel: this test, test_failed_validation_leaves_the_environment_alone and test_agrees_with_the_compilation_context all pass under it. Pre-existing, surfaced by this PR changing the behaviour they describe: - jit/core.py's "defaults to gfx942" comment and amd-flashinfer-jit-cache's README both still documented the hard-coded default this PR removes. - rocminfo_gpu_agents caught only FileNotFoundError and TimeoutExpired, so a present-but-unexecutable rocminfo (PermissionError) aborted resolution instead of reaching the documented GPU-less fallback. The removed _auto_detect_archs() swallowed these; catching OSError restores that. The eighth -- the module-global JIT compilation context being a different object from the one validated here -- is deferred with the other two design-level findings; see the thread replies.
|
Review 4985855053 — 8 suppressed; 7 fixed in b3fd023, 1 deferred. Four were defects the previous commit introduced. Adding Three were pre-existing and surfaced because this PR changes what they describe: Deferred: the module-global JIT context ( 164 passed; |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (3)
flashinfer/hip_utils.py:76
- The visibility resolver only reads
HIP_VISIBLE_DEVICESandCUDA_VISIBLE_DEVICES, but ROCr also appliesROCR_VISIBLE_DEVICESto HIP processes; this repository records that variable alongside the other GPU selectors inrocm_profiler.py:594-596. With onlyROCR_VISIBLE_DEVICESset on a mixed-architecture host, this code returns every physical agent and can compile for GPUs hidden from the process. Include the ROCr selector with a defined precedence and the same index/UUID handling.
raw = os.environ.get("HIP_VISIBLE_DEVICES")
if raw is None:
raw = os.environ.get("CUDA_VISIBLE_DEVICES")
if raw is None:
return agents
flashinfer/hip_utils.py:533
- The new partial-PyTorch filtering branch is still untested. The AOT test uses
gfx900,gfx950on ROCm 7.1, so step 1 removesgfx900before step 3, and the fake PyTorch flags include the remaininggfx950;missing_in_pytorchis therefore empty. Use two FlashInfer-supported targets (for examplegfx942,gfx950) with_get_rocm_arch_flags()returning onlygfx950, then assert the resultinggfx950list so this branch cannot regress unnoticed.
if missing_in_pytorch:
# Filter, then raise only if nothing is left -- matching steps 1 and
# 2, which already drop what they cannot serve and continue. Step 3
# raising outright made the GPU-less fallback unusable: it returns
# every supported architecture, so an arch-specific PyTorch wheel
flashinfer/hip_utils.py:80
-1is the CUDA/HIP no-device sentinel, but it failsisdigit()here and therefore takes the unmappable-device branch that returns all agents. A build launched withHIP_VISIBLE_DEVICES=-1orCUDA_VISIBLE_DEVICES=-1will consequently target hidden GPUs instead of taking the no-device fallback. Handle this sentinel as an empty visible set.
tokens = [t.strip() for t in raw.split(",") if t.strip()]
…step-3 filter Both found by self-review of the diff rather than by the review bot. - resolve_target_archs' docstring still described step 3 as "the supported GPUs actually present". Since visible_gpu_agents landed it is the supported GPUs *visible to this process*, which is the whole point of that change. - The previous commit changed the PyTorch check from raise-on-any-missing to filter-and-warn without adding a test. The existing raise test still passes because it requests a single arch that is entirely absent, so it never exercises the new branch. Added the partial case, A/B-verified against c3c9fb1~1 -- without the filter it fails with "Emitted warnings: []".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (3)
flashinfer/hip_utils.py:75
ROCR_VISIBLE_DEVICESis another ROCr visibility control used by HIP applications. If a launcher sets it (for example, to select only device 1) without settingHIP_VISIBLE_DEVICESorCUDA_VISIBLE_DEVICES, this path treats visibility as unset and returns everyrocminfoagent, so a mixed gfx942/gfx950 host can resolve an architecture that the process cannot access. IncludeROCR_VISIBLE_DEVICESin the precedence chain before the CUDA fallback, preserving an explicitly empty value.
raw = os.environ.get("HIP_VISIBLE_DEVICES")
if raw is None:
raw = os.environ.get("CUDA_VISIBLE_DEVICES")
if raw is None:
flashinfer/hip_utils.py:81
-1is the conventional HIP/CUDA sentinel for hiding all GPUs, but it is non-numeric here, so thenot all(...isdigit())branch returns every physical agent. A GPU-less build launched withHIP_VISIBLE_DEVICES=-1therefore skips the no-device fallback and can target hardware that the runtime exposes as unavailable. Treat the exact sentinel as an empty visibility selection and cover it with a regression test.
tokens = [t.strip() for t in raw.split(",") if t.strip()]
if not all(t.isdigit() for t in tokens):
flashinfer/hip_utils.py:537
- When
visible_gpu_agents()cannot map a UUID, it intentionally returns every physical architecture. This new filtering can then choose the wrong one: a process pinned to a gfx942 UUID on a mixed gfx942/gfx950 host with a gfx950-only PyTorch reportsgfx950as usable and silently builds for a GPU the process cannot see. Preserve the ambiguous-visibility state and fail (or require a fat PyTorch build) instead of applying this partial-support fallback to it.
usable = [f for f in arch_flags if f in pytorch_arch_flags]
…sentinel - The PyTorch filter added in c3c9fb1 applied to every request, so an explicit gfx942,gfx950 against a gfx942-only wheel quietly succeeded as gfx942 instead of failing. That is worse than the bug it fixed: the artifact no longer satisfies the target it claims. resolve_target_archs_with_origin() now reports explicit/env/detected/fallback, and only "fallback" -- a guess this module made with no hardware to look at -- may be narrowed. resolve_target_archs keeps its signature and delegates. - visible_gpu_agents only read HIP_/CUDA_VISIBLE_DEVICES. ROCR_VISIBLE_DEVICES also applies, and applies *beneath* HIP, so the two compose in a way indices cannot express: honoured alone, declined when combined. - "-1" is the no-device sentinel and is not a digit, so it took the unmappable branch and returned every agent -- the exact opposite of what it asks for. The test added in fdb8f38 used an explicit list and would now be wrong; it is replaced by a pair pinning both sides -- fallback narrows with a warning, explicit still raises.
|
Review 4987723853 — 3 suppressed, all already fixed in e512b84 (which landed after that review ran). ROCR_VISIBLE_DEVICES and the |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.
Suppressed comments (2)
flashinfer/hip_utils.py:100
- The
orchain does not actually give HIP visibility precedence whenHIP_VISIBLE_DEVICESis set to an empty/whitespace value (or-1): it falls through toCUDA_VISIBLE_DEVICES. For example, HIP="", CUDA="0"resolves GPU 0 instead of no devices, so the resolver can target hardware hidden from the HIP process. Select the variable by key presence so an explicitly empty HIP value remains authoritative.
raw = (
present.get("HIP_VISIBLE_DEVICES")
or present.get("ROCR_VISIBLE_DEVICES")
or present.get("CUDA_VISIBLE_DEVICES")
or ""
flashinfer/hip_utils.py:94
- The
ROCR_VISIBLE_DEVICEScomposition early return runs before the no-device handling below. Consequently,HIP_VISIBLE_DEVICES=-1or an explicitly empty HIP value combined with any ROCR value returns all physical agents, despite the documented no-device semantics; the resolver then treats the result as detected rather than as the GPU-less fallback and can reject an arch-specific PyTorch wheel for an invisible architecture. Handle deterministic no-device sentinels before declining to scope the composition.
if "ROCR_VISIBLE_DEVICES" in present and len(present) > 1:
logger.debug(
"Not scoping architecture detection: ROCR_VISIBLE_DEVICES composes "
"with %s and the result cannot be mapped to rocminfo's enumeration "
"order.",
", ".join(n for n in present if n != "ROCR_VISIBLE_DEVICES"),
)
return agents
…lity precedence The provenance added in e512b84 was discarded by both callers, so the fallback-only narrowing it gates never fired where it was needed. CompilationContext resolved to a string and passed it back in, which re-entered the resolver as "explicit"; aot_hip published to the environment first, which made it "env". A GPU-less host with an arch-specific wheel therefore still failed on the half it could not build. Verified end to end: that case now yields gfx942 with a warning, where it raised before. - CompilationContext passes arch_list=None and lets validation resolve, so the origin survives. It logs the validated set instead of the pre-resolved string. - aot_hip validates first and publishes after. Resolving twice is cheap now -- the probe is cached and the GPU-less warning is emitted once per message -- and publishing only on success also removes the need to restore on failure. Two precedence bugs in visible_gpu_agents, both from using `or` over values that are legitimately falsy: - HIP_VISIBLE_DEVICES="" means no devices, but empty is falsy, so the chain fell through to CUDA_VISIBLE_DEVICES and resolved GPUs the operator had hidden. Selection is now by key, in strict precedence order. - HIP_VISIBLE_DEVICES=-1 combined with any ROCR value hit the composition early return and got every agent. "No devices" is unambiguous whatever else is set, so it is now decided before the composition check. test_failed_validation_leaves_the_environment_alone asserted the old publish-before-validate order; it now asserts the reverse, which is the property that keeps provenance intact.
|
Review 4987846618 — 2 suppressed, both fixed in a421230. Both were mine, and both from using |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (2)
flashinfer/hip_utils.py:110
- The no-device check only examines the precedence-selected
rawvalue. If a composing layer is the no-device sentinel—for exampleHIP_VISIBLE_DEVICES=0withROCR_VISIBLE_DEVICES=-1or an empty ROCR value—this branch returns all physical agents instead of(), despite the documented empty/-1semantics. On a mixed-architecture host that can select invisible architectures and then fail PyTorch validation; handle no-device sentinels in the composing visibility variables before treating the composition as unknowable.
if "ROCR_VISIBLE_DEVICES" in present and len(present) > 1:
logger.debug(
"Not scoping architecture detection: ROCR_VISIBLE_DEVICES composes "
"with %s and the result cannot be mapped to rocminfo's enumeration "
"order.",
flashinfer/hip_utils.py:99
- A malformed signed token such as
HIP_VISIBLE_DEVICES=--1reaches this predicate:lstrip("+-").isdigit()is true, butint(t)is applied to the original--1and raisesValueError. Visibility input should degrade to the existing unmappable-device fallback rather than abort architecture resolution; only convert tokens after validating a single optional sign (or catch the conversion).
if not tokens or any(t.lstrip("+-").isdigit() and int(t) < 0 for t in tokens):
Both from Copilot review 4987970293, both mine.
- `HIP_VISIBLE_DEVICES=--1` crashed. `"--1".lstrip("+-").isdigit()` is True, so
it passed the sentinel guard and then raised ValueError inside int(). Parsing
is now total via a helper returning None, and anything unparseable takes the
existing unmappable-input path. Visibility comes off a launcher command line,
so it has to degrade rather than raise.
- The no-device check only looked at the precedence-selected variable, so
ROCR_VISIBLE_DEVICES=-1 alongside HIP_VISIBLE_DEVICES=0 fell through to the
composition branch and returned every agent, despite ROCr having hidden them
all. It now checks the selected variable and ROCR whenever ROCR is present,
since ROCr composes beneath HIP rather than being overridden by it. CUDA is
still only checked when selected, because HIP overrides it outright.
5 tests. The composing-layer one was vacuous on first write -- with a two-agent
list both behaviours return gfx942,gfx950 -- and only distinguishes them with a
single-agent list, where ignoring ROCR yields gfx942 and honouring it reaches
the fat fallback. Verified against a421230: 3 malformed cases and the
composing-layer case all fail there.
|
Review 4987970293 — 2 suppressed, both fixed in the push above. Both mine.
5 tests, A/B-verified against a421230. One of them was vacuous on first write and only distinguishes the two behaviours with a single-agent list. |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (2)
flashinfer/hip_utils.py:75
rocminfois an HSA/ROCr client, so unlikeHIP_VISIBLE_DEVICESit is affected byROCR_VISIBLE_DEVICES. WithROCR_VISIBLE_DEVICES=1,rocminfo_gpu_agents()can already return a one-element list for physical GPU 1, but this code then indexes that filtered list with1, drops the only agent, and falls back to building for all architectures. Collect the rocminfo inventory with the ROCR visibility variable removed (or otherwise distinguish an already-filtered inventory) before applying the requested indices.
agents = rocminfo_gpu_agents()
# ROCR_VISIBLE_DEVICES is applied by ROCr *beneath* HIP, so when it is set
# alongside one of the others the two compose and the composition cannot be
# reconstructed from indices alone. Refuse to scope in that case rather than
# apply the wrong one -- over-building is recoverable.
flashinfer/hip_utils.py:67
- Returning the full agent list for an unscopable UUID (and likewise a composed ROCR/HIP visibility setting) is not recoverable in the build path as written.
resolve_target_archs_with_origin()labels this result asdetected, and the PyTorch check invalidate_flashinfer_rocm_arch()only narrowsorigin == "fallback"; on a mixed gfx942/gfx950 host with a UUID selecting gfx950 and a gfx950-only PyTorch wheel, this therefore resolves both architectures and raises for missing gfx942 instead of building the visible target. Either map the UUID/composition, or carry an explicit “unscoped detection” provenance through validation so the conservative result can be narrowed rather than hard-failing.
from unset. A value naming devices by UUID (``GPU-...``) cannot be mapped to
an enumeration index here, so the full list is returned rather than guessing
-- over-building is recoverable, silently targeting the wrong card is not.
My ROCR handling was wrong in concept, not just in detail. rocminfo is an HSA
client, so ROCr has already filtered its output; re-applying ROCR_VISIBLE_DEVICES
indices on top filtered twice, and the "composition is ambiguous" case I invented
to work around it does not exist. Measured on this host:
rocminfo -> 1 GPU agent
ROCR_VISIBLE_DEVICES=-1 rocminfo -> 0 GPU agents (ROCr scopes it)
HIP_VISIBLE_DEVICES=-1 rocminfo -> 1 GPU agent (HIP does not)
That asymmetry is the entire reason this function exists, and it also makes HIP
compose with ROCr for free: HIP indices index into the ROCr-visible set, which
is exactly what rocminfo reports. Only HIP, then CUDA, is applied here now, and
the special-casing is gone.
Unmappable visibility (a UUID) now reports "nothing known to be visible" rather
than the full agent list. Returning everything labelled it `detected`, which the
provenance rule forbids narrowing, so an arch-specific wheel became a hard
failure. Reporting nothing reaches the fat fallback, which may be narrowed --
and warns, since a fat build is not what the operator asked for.
Two self-inflicted test failures on the way, both caught before pushing: a stale
test from the old model survived my replacement block, and removing it took the
parametrize decorator off the neighbouring test.
|
Review 4988146121 — 2 suppressed, both fixed. The first was a conceptual error on my part, not a detail.
Unmappable (UUID) visibility now reports "nothing known to be visible" instead of the full list: returning everything labelled it |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
flashinfer/aot_hip.py:249
- Validation success is not build success:
gen_all_modules()andbuild_jit_specs()below can still raise after this assignment. In that case the new value remains in the process-global environment (and can replace a caller's pre-existing value with a filtered list), contradicting the guarantee in lines 233-234 and influencing later JIT/AITER builds. Restore the prior environment value on downstream failure, or narrow that guarantee to validation failures.
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:618
- Because this probe is now a zero-argument process-wide cache, its result no longer tracks
ROCR_VISIBLE_DEVICES.visible_gpu_agents()deliberately leaves that variable to therocminfosubprocess, so if the first probe runs before that variable is set or after it changes, later resolution can reuse GPUs that ROCr has hidden. Key the cache by the ROCr visibility value (and any other environment that affectsrocminfo), or keep the raw probe uncached and cache only a visibility-aware result.
@functools.cache
| detected = sorted( | ||
| { | ||
| arch | ||
| for arch, _ in visible_gpu_agents() |
Summary
Three call sites independently answered "what architecture are we building for", and on CDNA4 they disagreed. Measured on a gfx950 host with
FLASHINFER_ROCM_ARCH_LISTunset, before this change:The JIT validated
gfx942while compiling forgfx950. That check exists to catch "your PyTorch was not built for this architecture", and it was asking about an architecture nobody was building for — vacuous on a PyTorch carrying both, and a spurious hard failure on an arch-specific build carrying only gfx950. Nothing noticed becausejit/core.pydiscards the return value; the comment onjit/core.py:148even said# ... or defaults to gfx942.What changed
hip_utils.resolve_target_archs()becomes the single resolver — explicit argument, thenFLASHINFER_ROCM_ARCH_LIST, then the architectures actually present, then every supported architecture with a warning.validate_rocm_arch,validate_flashinfer_rocm_arch,CompilationContextandaot_hipall route through it.CompilationContext._auto_detect_archsis deleted — private, no other caller, and the holder of two morereturn "gfx942"lines.The last-resort fallback changes from
"gfx942"to every supported architecture. On a GPU-less build host the old literal was not conservative, it was a guess that silently produced a gfx942-only artifact. A fat build is slower but correct wherever it lands, and the warning tells the operator how to make it cheap again. No Dockerfile, CI job or Jenkinsfile in this repo sets the variable, so this path is reachable in practice.Detection uses
rocminforather thantorch.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).The test that protected the bug
test_defaults_to_gfx942_when_no_env_and_no_argasserted== "gfx942"for exactly this path. A test that pins a wrong constant is how the constant survives. It now pins the detected architecture, and a newTestResolveTargetArchscovers all four resolution branches plus the property the change exists for — that the validator and the thing emitting--offload-archagree.Test plan
gfx950.FLASHINFER_ROCM_ARCH_LIST=gfx942on a gfx950 box still resolves togfx942— cross-compiling stays possible./dev/kfdresolves togfx942,gfx950, emits the warning, and does so without importing torch.test_hip_utils.py,test_aiter_build_arch_hip.py,test_arch_caps_hip.pyon gfx950.pre-commit runon all changed files.Deliberately not in this PR
aot_hipstill publishes the resolved list back intoos.environ, a process-global side effect that outlives the call. It is load-bearing rather than incidental: the AITER shim resolves its own build architecture fromFLASHINFER_ROCM_ARCH_LIST, and an AOT build has no other channel to tell it what this build targets. Removing it means threading an explicit parameter across the AOT → JIT boundary, which is wider than the resolution bug being fixed here. The value it publishes now comes from the same resolver, and the reasoning is recorded at the call site.Separately noted while measuring, not addressed:
import flashinferrequires a live GPU, not merely a ROCm-enabled torch —jit/env.py:212callstorch.cuda.get_device_properties(torch.cuda.current_device())at import time, which raisesNo HIP GPUs are availableon a build host.