[Feature] Support environments built by TheRock build system by eppaneamd · Pull Request #172 · AMD-Ecosystem/flashinfer

eppaneamd · 2026-02-20T13:55:58Z

📌 Description

This PR introduces ROCm/HIP compatibility improvements:

Refactor hip_utils.py for supporting environments built by TheRock build system
Other minor fixes

OBS!

Requires export PYTORCH_AMDCLANG=<path-to-amdclang++>.
Depends on [Fix] Add support for both CUDA and HIP generator headers in sampling.cu #180.

🧪 Tests

Currently failing tests:

=========================== short test summary info ============================
FAILED tests/test_logits_processor_hip.py::TestLogitsPipeCompilationHIP::test_probs_sample_freq[0.0-normal_distribution(std=1)-32000] - AssertionError: Compiled similarity: tensor([0.9719], device='cuda:0')
assert tensor([0.9719], device='cuda:0') > 0.99
FAILED tests/test_logits_processor_hip.py::TestLogitsPipeCompilationHIP::test_probs_sample_freq[0.0-normal_distribution(std=1)-128256] - AssertionError: Compiled similarity: tensor([0.9895], device='cuda:0')
assert tensor([0.9895], device='cuda:0') > 0.99
============ 2 failed, 589 passed, 12 skipped in 229.71s (0:03:49) =============

=========================== short test summary info ============================
FAILED tests/test_sampling_hip.py::test_sampling_freq[0.0-normal_distribution(std=1)-32000] - AssertionError: similarity: tensor([0.9725], device='cuda:0')
assert tensor([0.9725], device='cuda:0') > 0.99
FAILED tests/test_sampling_hip.py::test_sampling_freq[0.0-normal_distribution(std=1)-128256] - AssertionError: similarity: tensor([0.9894], device='cuda:0')
assert tensor([0.9894], device='cuda:0') > 0.99
============ 2 failed, 772 passed, 18 skipped in 283.44s (0:04:43) =============

Above similarity scores fall just below the 0.99 threshold (< 2% deviation).

Test environment:

OS: Ubuntu 22.04.5 LTS (x86_64)
GPU: AMD Instinct MI300X (gfx942:sramecc+:xnack-)
Python version: 3.11.10

rocm                            7.12.0a20260211
rocm-sdk-core                   7.12.0a20260211
rocm-sdk-libraries-gfx94X-dcgpu 7.12.0a20260211

HIP runtime version: 7.3.53390
ROCM used to build PyTorch: 7.3.53390

[pip3] numpy==2.3.1
[pip3] torch==2.10.0+rocm7.12.0a20260211
[pip3] torchao==0.16.0+git3c1065ca6
[pip3] torchaudio==2.10.0+rocm7.12.0a20260211
[pip3] torchcodec==0.10.0
[pip3] torchvision==0.25.0+rocm7.12.0a20260211
[pip3] triton==3.6.0+rocm7.12.0a20260211
[pip3] amd-aiter==0.1.10.post4.dev13+g439606621

…ader

This reverts commit 1909936.

Copilot

Pull request overview

This PR improves ROCm/HIP compatibility (including TheRock-built environments) by refactoring ROCm environment detection, making CUDA-specific imports conditional, and improving HIP compilation/header compatibility, while also adding a test-time GPU memory cleanup fixture.

Changes:

Refactor ROCm detection utilities to support TheRock layouts and centralize ROCM_HOME resolution.
Make CUDA-only imports in prefill.py conditional to avoid import-time failures on ROCm/CPU-only builds.
Update ROCm sampling extension to include the appropriate PyTorch generator header based on availability.
Add an autouse pytest fixture to opportunistically clear GPU memory between tests.

Reviewed changes

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

Show a summary per file

File	Description
tests/conftest.py	Adds an autouse GPU memory cleanup fixture and adjusts cache-clearing threshold logic.
flashinfer/prefill.py	Gates CUDA-only imports behind `torch.version.cuda` to improve ROCm/CPU importability.
flashinfer/jit/cpp_ext_hip.py	Switches ROCm home discovery to `flashinfer.hip_utils.get_rocm_home()` and simplifies ninja generation.
flashinfer/hip_utils.py	Adds TheRock detection + multiple ROCm version detection methods; refactors compatibility matrix construction.
flashinfer/csrc_rocm/sampling.cu	Selects HIP vs CUDA generator header using `__has_include` for better cross-torch compatibility.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

Comments suppressed due to low confidence (1)

tests/conftest.py:232

pytest_runtest_call is implemented as a generator (uses yield) but the hook decorator no longer sets wrapper=True. Without wrapper=True, pytest will not treat this as a hook wrapper and this will break test execution/exception handling. Set wrapper=True again (you can keep tryfirst=True if needed).

@pytest.hookimpl(tryfirst=True)
def pytest_runtest_call(item):
    # skip OOM error and missing JIT cache errors
    try:
        yield
    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

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

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-02-24T20:05:30Z

    import os
-    import re
-    import subprocess

-    # Method 1: Try /opt/rocm/.info/version (most reliable)
-    rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
-    version_file = os.path.join(rocm_path, ".info", "version")
+    return os.environ.get("ROCM_PATH") or os.environ.get("ROCM_HOME") or "/opt/rocm"


get_rocm_home() falls back to /opt/rocm without attempting any discovery when neither ROCM_PATH nor ROCM_HOME is set. This can break environments where ROCm is installed in a non-standard prefix (including TheRock wheel installs) unless users also set one of those env vars. Consider adding an additional fallback that derives the prefix from PYTORCH_AMDCLANG (if set) and/or from a hipcc/amdclang++ found on PATH, before defaulting to /opt/rocm.

Copilot · 2026-02-24T20:05:31Z

+            return version
+        print(f"ROCm version not found using {method.__name__}. Trying next method...")
+
+    return None
+
+


get_system_rocm_version() prints a message for every failed detection method attempt. This function is part of core validation and may run in library contexts where stdout noise is undesirable (e.g., tests, JIT compilation). Please route this through the project's logger at DEBUG level or gate it behind an explicit verbose/env flag instead of unconditional print().

diptorupd

Thank you @eppaneamd !

eppaneamd added 9 commits February 20, 2026 08:22

refactor system rocm detection + add therock support

960425f

update hipgen header for torch2.10 (bwd compat TODO)

274f5cc

cuda-specific imports

8d63dba

use hip generator header if available, otherwise fall back to cuda he…

03b9ac2

…ader

fix rocm arch groups

871f978

hook clear cuda cache to pytests

6d2c37a

enable torch custom op and register fake

1909936

Revert "enable torch custom op and register fake"

d5e9691

This reverts commit 1909936.

flashinfer test memory theshold to 0.75

11d49d0

Copilot AI review requested due to automatic review settings February 20, 2026 13:55

Copilot started reviewing on behalf of eppaneamd February 20, 2026 13:57 View session

Copilot AI reviewed Feb 20, 2026

View reviewed changes

Comment thread tests/conftest.py Outdated

Comment thread tests/conftest.py Outdated

Comment thread flashinfer/prefill.py Outdated

Comment thread flashinfer/hip_utils.py

Comment thread flashinfer/hip_utils.py

diptorupd reviewed Feb 20, 2026

View reviewed changes

Comment thread flashinfer/prefill.py

diptorupd reviewed Feb 20, 2026

View reviewed changes

Comment thread flashinfer/hip_utils.py

diptorupd reviewed Feb 20, 2026

View reviewed changes

Comment thread flashinfer/hip_utils.py

diptorupd reviewed Feb 21, 2026

View reviewed changes

Comment thread tests/conftest.py

diptorupd reviewed Feb 21, 2026

View reviewed changes

Comment thread flashinfer/csrc_rocm/sampling.cu

eppaneamd added 2 commits February 23, 2026 20:25

Merge branch 'amd-integration' into feat/enable_therock

0bd52f3

rm (move) sampling.cu changes to separate branch

fd189e5

Copilot AI review requested due to automatic review settings February 23, 2026 20:38

Copilot started reviewing on behalf of eppaneamd February 23, 2026 20:39 View session

Copilot AI reviewed Feb 23, 2026

View reviewed changes

Comment thread tests/conftest.py Outdated

eppaneamd added 2 commits February 23, 2026 21:22

revert prefill.py changes

66c6897

rm (move) conftest changes to another branch

d47465c

Copilot AI review requested due to automatic review settings February 23, 2026 21:32

Copilot started reviewing on behalf of eppaneamd February 23, 2026 21:34 View session

Copilot AI reviewed Feb 23, 2026

View reviewed changes

Comment thread flashinfer/hip_utils.py

eppaneamd and others added 2 commits February 24, 2026 08:07

detect therock build from rocm_sdk or therock manifest

a4e946c

Merge branch 'amd-integration' into feat/enable_therock

6b9be8d

Copilot AI review requested due to automatic review settings February 24, 2026 17:08

Copilot started reviewing on behalf of diptorupd February 24, 2026 17:10 View session

Copilot AI reviewed Feb 24, 2026

View reviewed changes

Comment thread flashinfer/hip_utils.py

Comment thread flashinfer/hip_utils.py

Comment thread flashinfer/hip_utils.py Outdated

Comment thread flashinfer/hip_utils.py Outdated

Comment thread flashinfer/hip_utils.py Outdated

eppaneamd added 2 commits February 24, 2026 19:58

fix docstring

1810bf5

fix docstring

4a70dff

Copilot AI review requested due to automatic review settings February 24, 2026 20:00

fix formatting

9aaec3c

Copilot started reviewing on behalf of eppaneamd February 24, 2026 20:02 View session

fix formatting

e2c89ff

Copilot AI reviewed Feb 24, 2026

View reviewed changes

diptorupd approved these changes Feb 24, 2026

View reviewed changes

diptorupd merged commit ff0328e into AMD-Ecosystem:amd-integration Feb 24, 2026
1 of 2 checks passed

Uh oh!

Conversation

eppaneamd commented Feb 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🧪 Tests

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 24, 2026

Choose a reason for hiding this comment

Uh oh!

diptorupd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

eppaneamd commented Feb 20, 2026 •

edited

Loading