From 960425f07c4a16611946c253da71cc4f1bd70d39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 08:22:31 +0000 Subject: [PATCH 01/17] refactor system rocm detection + add therock support --- flashinfer/hip_utils.py | 196 +++++++++++++++++++++++++--------- flashinfer/jit/cpp_ext_hip.py | 15 +-- 2 files changed, 151 insertions(+), 60 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index f676374057..5eee9673cb 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -6,40 +6,117 @@ FLASHINFER_SUPPORTED_ROCM_ARCHS = ["gfx942"] -def get_system_rocm_version(): +def get_rocm_home(): """ - Attempt to detect the system ROCm version. + Get the ROCM_HOME directory from environment variables or default path. Returns: - str: ROCm version like "6.4" or "7.0", or None if not detectable + str: Path to ROCm installation (e.g., "/opt/rocm") + """ + import os + + return os.environ.get("ROCM_PATH") or os.environ.get("ROCM_HOME") or "/opt/rocm" + + +def is_therock_build() -> bool: + """ + Check if ROCm was built by using TheRock build system. + + Returns: + bool: True if TheRock manifest exists, False otherwise + """ + import os + + rocm_home = get_rocm_home() + manifest_path = os.path.join(rocm_home, "share", "therock", "therock_manifest.json") + return os.path.isfile(manifest_path) + + +def get_system_rocm_version_from_info_file(): + """ + Try to get ROCm version from /opt/rocm/.info/version file. + + Returns: + str: ROCm version like "7.1.0" or None if not found """ 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") + rocm_home = get_rocm_home() + version_file = os.path.join(rocm_home, ".info", "version") try: with open(version_file, "r") as f: version = f.read().strip() return ".".join(version.split(".")[:3]) except (FileNotFoundError, IOError): + return None + + +def get_system_rocm_version_from_hipconfig(): + """ + Try to get ROCm version from hipconfig --version command. + + Returns: + str: ROCm version like "7.1.0" or None if not found + """ + import re + import subprocess + + try: + result = subprocess.run( + ["hipconfig", "--version"], + capture_output=True, + text=True, + timeout=5, + check=False + ) + if result.returncode == 0: + match = re.search(r"(\d+\.\d+\.\d+)", result.stdout) + if match: + return match.group(1) + except (subprocess.TimeoutExpired, FileNotFoundError): pass - # Method 2: Try amd-smi command + return None + + +def get_system_rocm_version_from_amd_smi(): + """ + Try to get ROCm version from amd-smi command. + + Returns: + str: ROCm version like "7.1.0" or None if not found + """ + import re + import subprocess + try: result = subprocess.run( - ["amd-smi"], capture_output=True, text=True, timeout=5, check=False + ["amd-smi"], + capture_output=True, + text=True, + timeout=5, + check=False ) if result.returncode == 0: - match = re.search(r"ROCm version:\s*(\d+\.\d+\.\d)", result.stdout) + match = re.search(r"ROCm version:\s*(\d+\.\d+\.\d+)", result.stdout) if match: return match.group(1) except (subprocess.TimeoutExpired, FileNotFoundError): pass - # Method 3: Try dpkg (Ubuntu/Debian) + return None + + +def get_system_rocm_version_from_dpkg(): + """ + Try to get ROCm version from dpkg (Ubuntu/Debian package manager). + + Returns: + str: ROCm version like "7.1.0" or None if not found + """ + import re + import subprocess + try: result = subprocess.run( ["dpkg", "-l", "rocm-core"], @@ -49,7 +126,7 @@ def get_system_rocm_version(): check=False, ) if result.returncode == 0: - match = re.search(r"rocm-core\s+(\d+\.\d+\.\d)", result.stdout) + match = re.search(r"rocm-core\s+(\d+\.\d+\.\d+)", result.stdout) if match: return match.group(1) except (subprocess.TimeoutExpired, FileNotFoundError): @@ -58,6 +135,37 @@ def get_system_rocm_version(): return None +def get_system_rocm_version(): + """ + Attempt to detect the system ROCm version. + + For standard builds, tries methods in order of reliability. + For TheRock builds, prioritizes hipconfig as it's more reliable. + + Returns: + str: ROCm version like "7.1.0" or None if not detectable + """ + # For TheRock builds, prioritize hipconfig + if is_therock_build(): + return get_system_rocm_version_from_hipconfig() + + # Try standard detection methods in order of reliability + detection_methods = [ + get_system_rocm_version_from_info_file, + get_system_rocm_version_from_amd_smi, + get_system_rocm_version_from_dpkg, + get_system_rocm_version_from_hipconfig, + ] + + for method in detection_methods: + version = method() + if version: + return version + print(f"ROCm version not found using {method.__name__}. Trying next method...") + + return None + + def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: """ Validate ROCm architecture against system ROCm version. @@ -77,42 +185,32 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: # ROCm compatibility matrix: version -> supported gfx architectures # Refer: https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html + # https://github.com/ROCm/TheRock/blob/main/SUPPORTED_GPUS.md#rocm-on-linux + # Update lists for adding or removing a version or arch + # Add new tuple for adding a new version group + _ROCM_ARCH_GROUPS = [ + ( + ["7.12", "7.11", "7.10"], + ["gfx950", "gfx942", "gfx90a", "gfx908", "gfx906", + "gfx1201", "gfx1200", "gfx1151", "gfx1150", + "gfx1102", "gfx1101", "gfx1100", "gfx1030"], + ), + ( + ["7.2", "7.1", "7.0"], + ["gfx950", "gfx1201", "gfx1200", "gfx1101", "gfx1100", + "gfx1030", "gfx942", "gfx90a", "gfx908"], + ), + ( + ["6.4", "6.3"], + ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"], + ), + ] + + # Build the compatibility matrix ROCM_COMPAT_MATRIX = { - "7.2": [ - "gfx950", - "gfx1201", - "gfx1200", - "gfx1101", - "gfx1100", - "gfx1030", - "gfx942", - "gfx90a", - "gfx908", - ], - "7.1": [ - "gfx950", - "gfx1201", - "gfx1200", - "gfx1101", - "gfx1100", - "gfx1030", - "gfx942", - "gfx90a", - "gfx908", - ], - "7.0": [ - "gfx950", - "gfx1201", - "gfx1200", - "gfx1101", - "gfx1100", - "gfx1030", - "gfx942", - "gfx90a", - "gfx908", - ], - "6.4": ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"], - "6.3": ["gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"], + version: archs + for versions, archs in _ROCM_ARCH_GROUPS + for version in versions } # Get architecture list from parameter, env var, or default @@ -124,7 +222,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: if system_rocm_version is None: raise RuntimeError( "Could not detect ROCm installation. Please ensure ROCm is installed and " - "accessible (check ROCM_PATH or /opt/rocm)." + "accessible (check ROCM_PATH, ROCM_HOME or /opt/rocm)." ) # Parse version to major.minor for compatibility check diff --git a/flashinfer/jit/cpp_ext_hip.py b/flashinfer/jit/cpp_ext_hip.py index 3b3ebad9dd..385013d080 100644 --- a/flashinfer/jit/cpp_ext_hip.py +++ b/flashinfer/jit/cpp_ext_hip.py @@ -19,10 +19,11 @@ _get_pybind11_abi_build_flags, ) -from torch.utils.cpp_extension import ROCM_HOME +from flashinfer.hip_utils import get_rocm_home from . import env as jit_env +ROCM_HOME = get_rocm_home() def _get_glibcxx_abi_build_flags() -> List[str]: glibcxx_abi_cflags = [ @@ -112,11 +113,7 @@ def generate_ninja_build_for_op( ldflags += extra_ldflags cxx = os.environ.get("CXX", "c++") - rocm_home = ROCM_HOME or "/opt/rocm" - amdclang = os.environ.get("PYTORCH_AMDCLANG", "$rocm_home/bin/amdclang++") - - cxx = os.environ.get("CXX", "c++") - rocm_home = ROCM_HOME or "/opt/rocm" + rocm_home = ROCM_HOME amdclang = os.environ.get("PYTORCH_AMDCLANG", "$rocm_home/bin/amdclang++") lines = [ @@ -164,11 +161,7 @@ def generate_ninja_build_for_op( for source in sources: is_hip = source.suffix == ".cu" object_suffix = ".cuda.o" if is_hip else ".o" - cmd = "" - if is_hip: - cmd = "hip_compile" - else: - cmd = "compile" + cmd = "hip_compile" if is_hip else "compile" obj_name = source.with_suffix(object_suffix).name obj = f"$name/{obj_name}" objects.append(obj) From 274f5cc3f1d82bd15986e8d74ba9217c7b857fcf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 08:24:06 +0000 Subject: [PATCH 02/17] update hipgen header for torch2.10 (bwd compat TODO) --- flashinfer/csrc_rocm/sampling.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashinfer/csrc_rocm/sampling.cu b/flashinfer/csrc_rocm/sampling.cu index c1afc2019a..dd486764b7 100644 --- a/flashinfer/csrc_rocm/sampling.cu +++ b/flashinfer/csrc_rocm/sampling.cu @@ -11,7 +11,7 @@ typedef hipStream_t cudaStream_t; #include #include -#include +#include #include #include From 8d63dba39ff4ba5c7d4ba069c54c272d60f6a37f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 09:46:54 +0000 Subject: [PATCH 03/17] cuda-specific imports --- flashinfer/prefill.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 420ce4b53b..3f58be222b 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -25,14 +25,20 @@ from .jit import ( gen_batch_prefill_module, gen_customize_batch_prefill_module, - gen_fmha_cutlass_sm100a_module, gen_single_prefill_module, get_batch_prefill_uri, get_single_prefill_uri, - setup_cubin_loader, - gen_trtllm_gen_fmha_module, ) -from .cudnn import cudnn_batch_prefill_with_kv_cache + +# CUDA-specific imports +if torch.version.cuda: + from .jit import ( + gen_fmha_cutlass_sm100a_module, + setup_cubin_loader, + gen_trtllm_gen_fmha_module, + ) + from .cudnn import cudnn_batch_prefill_with_kv_cache + from .page import block_sparse_indices_to_vector_sparse_offsets, get_seq_lens from .quantization import packbits, segment_packbits from .utils import ( From 03b9ac2ef613365e120ca11e4c342ab72c26612b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 09:58:23 +0000 Subject: [PATCH 04/17] use hip generator header if available, otherwise fall back to cuda header --- flashinfer/csrc_rocm/sampling.cu | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/flashinfer/csrc_rocm/sampling.cu b/flashinfer/csrc_rocm/sampling.cu index dd486764b7..405bfe34fb 100644 --- a/flashinfer/csrc_rocm/sampling.cu +++ b/flashinfer/csrc_rocm/sampling.cu @@ -11,7 +11,13 @@ typedef hipStream_t cudaStream_t; #include #include + +// Use HIP generator header if available, otherwise fall back to CUDA header +#if __has_include() #include +#else +#include +#endif #include #include From 871f97897020a752f706a65f873bd120aca3ac10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 10:13:01 +0000 Subject: [PATCH 05/17] fix rocm arch groups --- flashinfer/hip_utils.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 5eee9673cb..b036fad6f1 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -190,13 +190,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: # Add new tuple for adding a new version group _ROCM_ARCH_GROUPS = [ ( - ["7.12", "7.11", "7.10"], - ["gfx950", "gfx942", "gfx90a", "gfx908", "gfx906", - "gfx1201", "gfx1200", "gfx1151", "gfx1150", - "gfx1102", "gfx1101", "gfx1100", "gfx1030"], - ), - ( - ["7.2", "7.1", "7.0"], + ["7.3", "7.2", "7.1", "7.0"], ["gfx950", "gfx1201", "gfx1200", "gfx1101", "gfx1100", "gfx1030", "gfx942", "gfx90a", "gfx908"], ), From 6d2c37a9965e58197dba4cb65d9aad9a31f704a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 12:03:33 +0000 Subject: [PATCH 06/17] hook clear cuda cache to pytests --- tests/conftest.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 38da728c02..0bc8754057 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -195,6 +195,26 @@ def is_cuda_oom_error_str(e: str) -> bool: return "CUDA" in e and "out of memory" in e +def clear_cuda_cache(device: torch.device) -> None: + total_memory = get_device_properties(device).total_memory + reserved_memory = torch.cuda.memory_reserved() + + # FLASHINFER_TEST_MEMORY_THRESHOLD: threshold for PyTorch reserved memory usage (default: 0.9) + threshold = float(os.environ.get("FLASHINFER_TEST_MEMORY_THRESHOLD", "0.9")) + + if reserved_memory > threshold * total_memory: + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture(autouse=True, scope="function") +def clear_gpu_memory(): + yield + if torch.cuda.is_available(): + device = torch.device("cuda:0") + clear_cuda_cache(device) # Use the existing function + + @pytest.hookimpl(tryfirst=True) def pytest_runtest_call(item): # skip OOM error @@ -212,18 +232,6 @@ def get_device_properties(device: torch.device): return torch.cuda.get_device_properties(device) -def clear_cuda_cache(device: torch.device) -> None: - total_memory = get_device_properties(device).total_memory - reserved_memory = torch.cuda.memory_reserved() - - # FLASHINFER_TEST_MEMORY_THRESHOLD: threshold for PyTorch reserved memory usage (default: 0.9) - threshold = float(os.environ.get("FLASHINFER_TEST_MEMORY_THRESHOLD", "0.9")) - - if reserved_memory > threshold * total_memory: - gc.collect() - torch.cuda.empty_cache() - - # collected from gsk8k trace in sglang VARLEN_INDPTR_PARAMS = [ [ From 1909936d2b9774345f3fc766dd67e3e7c82e26f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 12:18:17 +0000 Subject: [PATCH 07/17] enable torch custom op and register fake --- flashinfer/utils.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/flashinfer/utils.py b/flashinfer/utils.py index 263e23af67..4228a8fbbf 100644 --- a/flashinfer/utils.py +++ b/flashinfer/utils.py @@ -272,22 +272,21 @@ def register_custom_op( ) -> Callable: # NOTE(Zihao): torch.library.custom_op has significant overhead as mentioned in the following link # https://github.com/vllm-project/vllm/blob/36e76700453924c8d421db99af70a88a1df835cd/vllm/utils.py#L1660-L1674 - - # return torch.library.custom_op( - # name, - # fn, - # mutates_args=mutates_args, - # device_types=device_types, - # schema=schema, - # ) - return lambda x: x + # TODO: doublecheck if this overhead is still significant and if so, consider implementing a caching mechanism to mitigate it + + return torch.library.custom_op( + name, + fn, + mutates_args=mutates_args, + device_types=device_types, + schema=schema, + ) def register_fake_op( name: str, fn: Optional[Callable] = None, ) -> Callable: - # return torch.library.register_fake(name, fn) - return lambda x: x + return torch.library.register_fake(name, fn) def determine_gemm_backend(device: torch.device) -> str: From d5e96912f0af1b11abc813f91ab95e25a84b4c34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 12:34:29 +0000 Subject: [PATCH 08/17] Revert "enable torch custom op and register fake" This reverts commit 1909936d2b9774345f3fc766dd67e3e7c82e26f3. --- flashinfer/utils.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/flashinfer/utils.py b/flashinfer/utils.py index 4228a8fbbf..263e23af67 100644 --- a/flashinfer/utils.py +++ b/flashinfer/utils.py @@ -272,21 +272,22 @@ def register_custom_op( ) -> Callable: # NOTE(Zihao): torch.library.custom_op has significant overhead as mentioned in the following link # https://github.com/vllm-project/vllm/blob/36e76700453924c8d421db99af70a88a1df835cd/vllm/utils.py#L1660-L1674 - # TODO: doublecheck if this overhead is still significant and if so, consider implementing a caching mechanism to mitigate it - - return torch.library.custom_op( - name, - fn, - mutates_args=mutates_args, - device_types=device_types, - schema=schema, - ) + + # return torch.library.custom_op( + # name, + # fn, + # mutates_args=mutates_args, + # device_types=device_types, + # schema=schema, + # ) + return lambda x: x def register_fake_op( name: str, fn: Optional[Callable] = None, ) -> Callable: - return torch.library.register_fake(name, fn) + # return torch.library.register_fake(name, fn) + return lambda x: x def determine_gemm_backend(device: torch.device) -> str: From 11d49d0f908810c535c1022fce13dcb14ac7ad15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Fri, 20 Feb 2026 13:24:45 +0000 Subject: [PATCH 09/17] flashinfer test memory theshold to 0.75 --- tests/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 0bc8754057..e15ce1b989 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -199,8 +199,8 @@ def clear_cuda_cache(device: torch.device) -> None: total_memory = get_device_properties(device).total_memory reserved_memory = torch.cuda.memory_reserved() - # FLASHINFER_TEST_MEMORY_THRESHOLD: threshold for PyTorch reserved memory usage (default: 0.9) - threshold = float(os.environ.get("FLASHINFER_TEST_MEMORY_THRESHOLD", "0.9")) + # FLASHINFER_TEST_MEMORY_THRESHOLD: threshold for PyTorch reserved memory usage (default: 0.75) + threshold = float(os.environ.get("FLASHINFER_TEST_MEMORY_THRESHOLD", "0.75")) if reserved_memory > threshold * total_memory: gc.collect() From fd189e5e1a7d97fb62ae549b5c3d1623a6868804 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Mon, 23 Feb 2026 20:38:07 +0000 Subject: [PATCH 10/17] rm (move) sampling.cu changes to separate branch --- flashinfer/csrc_rocm/sampling.cu | 8 -------- 1 file changed, 8 deletions(-) diff --git a/flashinfer/csrc_rocm/sampling.cu b/flashinfer/csrc_rocm/sampling.cu index ccbce4e250..a0b94bf7ee 100644 --- a/flashinfer/csrc_rocm/sampling.cu +++ b/flashinfer/csrc_rocm/sampling.cu @@ -10,14 +10,6 @@ typedef hipStream_t cudaStream_t; #endif #include -#include - -// Use HIP generator header if available, otherwise fall back to CUDA header -#if __has_include() -#include -#else -#include -#endif #include From 66c6897420d87f834f75aeb05adbfbf2895bd991 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Mon, 23 Feb 2026 21:22:59 +0000 Subject: [PATCH 11/17] revert prefill.py changes --- flashinfer/prefill.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/flashinfer/prefill.py b/flashinfer/prefill.py index 4fca3079ee..47d725c5d3 100755 --- a/flashinfer/prefill.py +++ b/flashinfer/prefill.py @@ -25,20 +25,14 @@ from .jit import ( gen_batch_prefill_module, gen_customize_batch_prefill_module, + gen_fmha_cutlass_sm100a_module, gen_single_prefill_module, get_batch_prefill_uri, get_single_prefill_uri, + setup_cubin_loader, + gen_trtllm_gen_fmha_module, ) - -# CUDA-specific imports -if torch.version.cuda: - from .jit import ( - gen_fmha_cutlass_sm100a_module, - setup_cubin_loader, - gen_trtllm_gen_fmha_module, - ) - from .cudnn import cudnn_batch_prefill_with_kv_cache - +from .cudnn import cudnn_batch_prefill_with_kv_cache from .page import block_sparse_indices_to_vector_sparse_offsets, get_seq_lens from .quantization import packbits, segment_packbits from .utils import ( From d47465c1768716805f14560dc2f4df053693e339 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Mon, 23 Feb 2026 21:32:18 +0000 Subject: [PATCH 12/17] rm (move) conftest changes to another branch --- tests/conftest.py | 23 +---------------------- 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 4261e71f1d..7bb16e204b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,3 @@ -import gc import json import os import types @@ -204,27 +203,7 @@ def is_cuda_oom_error_str(e: str) -> bool: return "CUDA" in e and "out of memory" in e -def clear_cuda_cache(device: torch.device) -> None: - total_memory = get_device_properties(device).total_memory - reserved_memory = torch.cuda.memory_reserved() - - # FLASHINFER_TEST_MEMORY_THRESHOLD: threshold for PyTorch reserved memory usage (default: 0.75) - threshold = float(os.environ.get("FLASHINFER_TEST_MEMORY_THRESHOLD", "0.75")) - - if reserved_memory > threshold * total_memory: - gc.collect() - torch.cuda.empty_cache() - - -@pytest.fixture(autouse=True, scope="function") -def clear_gpu_memory(): - yield - if torch.cuda.is_available(): - device = torch.device("cuda:0") - clear_cuda_cache(device) # Use the existing function - - -@pytest.hookimpl(tryfirst=True) +@pytest.hookimpl(wrapper=True) def pytest_runtest_call(item): # skip OOM error and missing JIT cache errors try: From a4e946c0fd6354d3d7335deb77c2fdf42e4797ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 24 Feb 2026 08:07:06 +0000 Subject: [PATCH 13/17] detect therock build from rocm_sdk or therock manifest --- flashinfer/hip_utils.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index cd5fa87e8f..02b4f1dc37 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -23,10 +23,19 @@ def is_therock_build() -> bool: Check if ROCm was built by using TheRock build system. Returns: - bool: True if TheRock manifest exists, False otherwise + bool: True if TheRock build is detected, False otherwise """ import os + # First, try checking for rocm_sdk package + try: + import rocm_sdk + if hasattr(rocm_sdk, '__version__') and rocm_sdk.__version__: + return True + except ImportError: + pass + + # Fall back to checking for TheRock manifest file rocm_home = get_rocm_home() manifest_path = os.path.join(rocm_home, "share", "therock", "therock_manifest.json") return os.path.isfile(manifest_path) From 1810bf56ddf11415c7f9f9dae981ded101f23103 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 24 Feb 2026 19:58:42 +0000 Subject: [PATCH 14/17] fix docstring --- flashinfer/hip_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 02b4f1dc37..2d3d1e4c28 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -43,7 +43,7 @@ def is_therock_build() -> bool: def get_system_rocm_version_from_info_file(): """ - Try to get ROCm version from /opt/rocm/.info/version file. + Try to get ROCm version from .info/version file located in ROCM_HOME. Returns: str: ROCm version like "7.1.0" or None if not found From 4a70dff4f25e077ff30440d761c96f22ddbb1365 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 24 Feb 2026 20:00:37 +0000 Subject: [PATCH 15/17] fix docstring --- flashinfer/hip_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 2d3d1e4c28..2570f5cab3 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -20,7 +20,7 @@ def get_rocm_home(): def is_therock_build() -> bool: """ - Check if ROCm was built by using TheRock build system. + Check if ROCm was built using TheRock build system. Returns: bool: True if TheRock build is detected, False otherwise From 9aaec3cee60b0e9625b6f3d575ae3075fed8d470 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 24 Feb 2026 20:01:21 +0000 Subject: [PATCH 16/17] fix formatting --- flashinfer/hip_utils.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index 2570f5cab3..3763498fd1 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -30,7 +30,8 @@ def is_therock_build() -> bool: # First, try checking for rocm_sdk package try: import rocm_sdk - if hasattr(rocm_sdk, '__version__') and rocm_sdk.__version__: + + if hasattr(rocm_sdk, "__version__") and rocm_sdk.__version__: return True except ImportError: pass @@ -76,7 +77,7 @@ def get_system_rocm_version_from_hipconfig(): capture_output=True, text=True, timeout=5, - check=False + check=False, ) if result.returncode == 0: match = re.search(r"(\d+\.\d+\.\d+)", result.stdout) @@ -100,11 +101,7 @@ def get_system_rocm_version_from_amd_smi(): try: result = subprocess.run( - ["amd-smi"], - capture_output=True, - text=True, - timeout=5, - check=False + ["amd-smi"], capture_output=True, text=True, timeout=5, check=False ) if result.returncode == 0: match = re.search(r"ROCm version:\s*(\d+\.\d+\.\d+)", result.stdout) @@ -200,8 +197,17 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: _ROCM_ARCH_GROUPS = [ ( ["7.3", "7.2", "7.1", "7.0"], - ["gfx950", "gfx1201", "gfx1200", "gfx1101", "gfx1100", - "gfx1030", "gfx942", "gfx90a", "gfx908"], + [ + "gfx950", + "gfx1201", + "gfx1200", + "gfx1101", + "gfx1100", + "gfx1030", + "gfx942", + "gfx90a", + "gfx908", + ], ), ( ["6.4", "6.3"], @@ -211,9 +217,7 @@ def validate_rocm_arch(arch_list: str = None, verbose: bool = False) -> str: # Build the compatibility matrix ROCM_COMPAT_MATRIX = { - version: archs - for versions, archs in _ROCM_ARCH_GROUPS - for version in versions + version: archs for versions, archs in _ROCM_ARCH_GROUPS for version in versions } # Get architecture list from parameter, env var, or default From e2c89ff4e65d479c229cdb8f65430d1a2e73a615 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rony=20Lepp=C3=A4nen?= Date: Tue, 24 Feb 2026 20:05:16 +0000 Subject: [PATCH 17/17] fix formatting --- flashinfer/jit/cpp_ext_hip.py | 1 + 1 file changed, 1 insertion(+) diff --git a/flashinfer/jit/cpp_ext_hip.py b/flashinfer/jit/cpp_ext_hip.py index 385013d080..f3df6a8a84 100644 --- a/flashinfer/jit/cpp_ext_hip.py +++ b/flashinfer/jit/cpp_ext_hip.py @@ -25,6 +25,7 @@ ROCM_HOME = get_rocm_home() + def _get_glibcxx_abi_build_flags() -> List[str]: glibcxx_abi_cflags = [ "-D_GLIBCXX_USE_CXX11_ABI=" + str(int(torch._C._GLIBCXX_USE_CXX11_ABI))