Detect ROCm wheel variant from environment for precompiled wheels. (#49365)
Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Aarushi Jain committed
Aug 13, 2026 at 23:19 UTC
59d1af569854a23f268fc7bfbf578fc06f9c166a
2 files changed
+302
-12
setup.py
+244
-10
@@ -507,20 +507,69 @@ class precompiled_wheel_utils:
507
508
@staticmethod
509
def fetch_metadata_for_variant(
510
- commit: str, variant: str | None
510
+ commit: str,
511
+ variant: str | None,
512
+ *,
513
+ rocm: bool = False,
514
) -> tuple[list[dict], str]:
515
"""
516
Fetches metadata for a specific variant of the precompiled wheel.
517
+
518
+ For non-ROCm, fetches vllm metadata.
519
+
520
+ For ROCm, discovers all first-level packages and combines their
521
+ metadata into a single list.
522
"""
515
- variant_dir = f"{variant}/" if variant is not None else ""
516
- repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
517
- meta_url = repo_url + "metadata.json"
518
- print(f"Trying to fetch nightly build metadata from {meta_url}")
523
+ import json
524
+ from html.parser import HTMLParser
525
from urllib.request import urlopen
526
521
- with urlopen(meta_url) as resp:
522
- # urlopen raises HTTPError on unexpected status code
523
- wheels = json.loads(resp.read().decode("utf-8"))
527
+ variant_dir = f"{variant}/" if variant is not None else ""
528
+
529
+ if not rocm:
530
+ # Keep original behavior
531
+ repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
532
+ meta_url = repo_url + "metadata.json"
533
+ print(f"Trying to fetch nightly build metadata from {meta_url}")
534
+ with urlopen(meta_url) as resp:
535
+ wheels = json.loads(resp.read().decode("utf-8"))
536
+
537
+ return wheels, repo_url
538
+
539
+ # ROCm: discover all packages under the variant directory.
540
+ repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}"
541
+
542
+ class LinkParser(HTMLParser):
543
+ def __init__(self):
544
+ super().__init__()
545
+ self.links: list[str] = []
546
+
547
+ def handle_starttag(self, tag, attrs):
548
+ if tag == "a":
549
+ href = dict(attrs).get("href")
550
+ if href:
551
+ self.links.append(href)
552
+
553
+ with urlopen(repo_url) as resp:
554
+ parser = LinkParser()
555
+ parser.feed(resp.read().decode("utf-8"))
556
+
557
+ packages = [
558
+ href.rstrip("/")
559
+ for href in parser.links
560
+ if href.endswith("/")
561
+ and href not in ("../", "./", "/")
562
+ and "/" not in href.rstrip("/")
563
+ ]
564
+
565
+ wheels: list[dict] = []
566
+
567
+ for package in packages:
568
+ meta_url = f"{repo_url}{package}/metadata.json"
569
+ print(f"Trying to fetch nightly build metadata from {meta_url}")
570
+ with urlopen(meta_url) as resp:
571
+ package_wheels = json.loads(resp.read().decode("utf-8"))
572
+ wheels.extend(package_wheels)
573
return wheels, repo_url
574
575
@staticmethod
@@ -582,6 +631,141 @@ class precompiled_wheel_utils:
631
print(f"Detected CUDA {cuda_version}, using variant {variant}")
632
return variant
633
634
+ @staticmethod
635
+ def rocm_version_to_variant(rocm_version: str) -> str:
636
+ """Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723."""
637
+ return "rocm" + rocm_version.replace(".", "")
638
+
639
+ @staticmethod
640
+ def detect_system_rocm_variant() -> str | None:
641
+ """Auto-detect the ROCm wheel variant from the installed ROCm stack."""
642
+ rocm_version = get_rocm_version()
643
+ if not rocm_version:
644
+ try:
645
+ import torch
646
+
647
+ rocm_version = torch.version.hip
648
+ except Exception:
649
+ pass
650
+ if not rocm_version:
651
+ return None
652
+ variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version)
653
+ print(f"Detected ROCm {rocm_version}, using variant {variant}")
654
+ return variant
655
+
656
+ @staticmethod
657
+ def fetch_available_rocm_variants(commit: str) -> list[str]:
658
+ """List ROCm wheel variants published for a commit on wheels.vllm.ai."""
659
+ from urllib.request import urlopen
660
+
661
+ index_url = f"https://wheels.vllm.ai/rocm/{commit}/"
662
+ print(f"Fetching available ROCm variants from {index_url}")
663
+ try:
664
+ with urlopen(index_url) as resp:
665
+ html = resp.read().decode("utf-8")
666
+ except Exception as e:
667
+ logger.warning(
668
+ "Failed to fetch ROCm variant index for commit %s: %s", commit, e
669
+ )
670
+ return []
671
+ variants = sorted(set(re.findall(r"rocm\d+", html)))
672
+ print(f"Available ROCm variants for commit {commit}: {variants}")
673
+ return variants
674
+
675
+ @staticmethod
676
+ def resolve_rocm_wheel_variant(
677
+ commit: str, variant_override: str | None
678
+ ) -> str | None:
679
+ """Resolve a ROCm wheel variant for a commit from wheels.vllm.ai."""
680
+ env_variant = precompiled_wheel_utils.detect_system_rocm_variant()
681
+ available = precompiled_wheel_utils.fetch_available_rocm_variants(commit)
682
+
683
+ if variant_override is not None:
684
+ if env_variant and variant_override != env_variant:
685
+ logger.warning(
686
+ "VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the "
687
+ "detected environment ROCm variant %s; refusing to use a "
688
+ "different ROCm patch wheel",
689
+ variant_override,
690
+ env_variant,
691
+ )
692
+ return None
693
+ if available and variant_override not in available:
694
+ logger.warning(
695
+ "Requested ROCm variant %s is not available for commit %s "
696
+ "(available: %s)",
697
+ variant_override,
698
+ commit,
699
+ available,
700
+ )
701
+ return None
702
+ return variant_override
703
+
704
+ if env_variant is None:
705
+ if available:
706
+ print(
707
+ "Could not detect ROCm version from the environment; "
708
+ f"available variants for commit {commit}: {available}"
709
+ )
710
+ return None
711
+
712
+ if env_variant in available:
713
+ return env_variant
714
+
715
+ logger.warning(
716
+ "Environment ROCm variant %s is not available for commit %s "
717
+ "(available: %s). Precompiled wheels may not be compatible.",
718
+ env_variant,
719
+ commit,
720
+ available,
721
+ )
722
+ return None
723
+
724
+ @staticmethod
725
+ def warn_if_rocm_torch_version_mismatch(
726
+ wheels: list[dict], repo_url: str, arch: str
727
+ ) -> None:
728
+ """Warn if installed torch differs from the custom ROCm build on
729
+ wheels.vllm.ai and suggest the correct install command."""
730
+ try:
731
+ installed = torch.__version__
732
+ except Exception:
733
+ return
734
+
735
+ def _wheel_version(pkg: str) -> str | None:
736
+ for w in wheels:
737
+ if w.get("package_name") == pkg and arch in w.get("platform_tag", ""):
738
+ v = w["version"]
739
+ if w.get("variant") and "+" not in v:
740
+ v = f"{v}+{w['variant']}"
741
+ return v
742
+ return None
743
+
744
+ expected = _wheel_version("torch")
745
+ if expected is None or installed == expected:
746
+ return
747
+
748
+ pkgs = f"torch=={expected}"
749
+ triton_ver = _wheel_version("triton")
750
+ if triton_ver:
751
+ pkgs += f" triton=={triton_ver}"
752
+
753
+ logger.warning(
754
+ "Installed PyTorch %s does not match the custom build %s "
755
+ "shipped with vLLM ROCm wheels. The ABI may differ from "
756
+ "official releases. If you hit extension load errors, "
757
+ "reinstall from the vLLM index:\n"
758
+ " pip install %s --extra-index-url %s\n"
759
+ " uv pip install %s --extra-index-url %s "
760
+ "--index-strategy unsafe-best-match",
761
+ installed,
762
+ expected,
763
+ pkgs,
764
+ repo_url,
765
+ pkgs,
766
+ repo_url,
767
+ )
768
+
769
@staticmethod
770
def find_local_rocm_wheel() -> str | None:
771
"""Search for a local vllm wheel in common locations."""
@@ -639,6 +823,57 @@ class precompiled_wheel_utils:
823
print(f"Found local ROCm wheel: {local_wheel}")
824
return local_wheel, None
825
826
+ import platform
827
+
828
+ arch = platform.machine()
829
+ commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
830
+ if not commit or len(commit) != 40:
831
+ print(
832
+ f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}"
833
+ ", trying to fetch base commit in main branch"
834
+ )
835
+ commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
836
+ variant = precompiled_wheel_utils.resolve_rocm_wheel_variant(
837
+ commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
838
+ )
839
+ print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}")
840
+ wheels, repo_url = None, None
841
+ if variant is not None:
842
+ try:
843
+ wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
844
+ commit, variant, rocm=True
845
+ )
846
+ precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch(
847
+ wheels, repo_url, arch
848
+ )
849
+ except Exception as e:
850
+ logger.warning(
851
+ "Failed to fetch ROCm wheel metadata for variant %s: %s",
852
+ variant,
853
+ e,
854
+ )
855
+ if wheels is not None and repo_url is not None:
856
+ from urllib.parse import urljoin
857
+
858
+ for wheel in wheels:
859
+ if wheel.get("package_name") == "vllm" and arch in wheel.get(
860
+ "platform_tag", ""
861
+ ):
862
+ print(f"Found precompiled wheel metadata: {wheel}")
863
+ if "path" not in wheel:
864
+ raise ValueError(f"Wheel metadata missing path: {wheel}")
865
+ wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"])
866
+ download_filename = wheel.get("filename")
867
+ print(f"Using precompiled wheel URL: {wheel_url}")
868
+ return wheel_url, download_filename
869
+ logger.warning(
870
+ "No precompiled vllm wheel found for architecture %s "
871
+ "from repo %s. All available wheels: %s",
872
+ arch,
873
+ repo_url,
874
+ wheels,
875
+ )
876
+
877
# Fall back to AMD's PyPI index
878
index_url = os.getenv(
879
"VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple"
@@ -669,8 +904,7 @@ class precompiled_wheel_utils:
904
print(f"Using user-specified precompiled wheel location: {wheel_location}")
905
return wheel_location, None
906
else:
672
- # ROCm: use local wheel or AMD's PyPI index
673
- # TODO: When we have ROCm nightly wheels, we can update this logic.
907
+ # ROCm: resolve wheels from wheels.vllm.ai with environment-matched variant
908
if precompiled_wheel_utils.is_rocm_system():
909
return precompiled_wheel_utils.determine_wheel_url_rocm()
910
tests/standalone_tests/python_only_compile.sh
+58
-2
@@ -28,8 +28,64 @@ fi
28
# test whether the metadata.json url is valid, retry each 3 minutes up to 5 times
29
# this avoids cumbersome error messages & manual retries in case the precompiled wheel
30
# for the given commit is still being built in the release pipeline
31
-meta_json_url="https://wheels.vllm.ai/$merge_base_commit/vllm/metadata.json"
32
-echo "INFO: will use metadata.json from $meta_json_url"
31
+_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')"
32
+if [[ "${_vllm_target_lower}" == "rocm" ]] || [[ -d /opt/rocm ]] || command -v rocminfo >/dev/null 2>&1; then
33
+ _rocm_env_variant="$(python3 - <<'PY'
34
+import ctypes
35
+import os
36
+from pathlib import Path
37
+
38
+def get_rocm_version() -> str | None:
39
+ rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") or "/opt/rocm"
40
+ try:
41
+ librocm_core = Path(rocm_home) / "lib" / "librocm-core.so"
42
+ if not librocm_core.is_file():
43
+ return None
44
+ librocm = ctypes.CDLL(str(librocm_core))
45
+ get_rocm_core_version = librocm.getROCmVersion
46
+ major = ctypes.c_uint32()
47
+ minor = ctypes.c_uint32()
48
+ patch = ctypes.c_uint32()
49
+ if get_rocm_core_version(
50
+ ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
51
+ ) == 0:
52
+ return f"{major.value}.{minor.value}.{patch.value}"
53
+ except Exception:
54
+ return None
55
+ return None
56
+
57
+version = get_rocm_version()
58
+if version:
59
+ print(f"rocm{version.replace('.', '')}", end="")
60
+PY
61
+)"
62
+ _available_variants="$(curl -sf "https://wheels.vllm.ai/rocm/${merge_base_commit}/" \
63
+ | grep -oP 'rocm\d+' | sort -u | tr '\n' ' ' || true)"
64
+ if [[ -n "${VLLM_PRECOMPILED_WHEEL_VARIANT:-}" ]]; then
65
+ _rocm_variant="${VLLM_PRECOMPILED_WHEEL_VARIANT}"
66
+ if [[ -n "${_rocm_env_variant}" && "${_rocm_variant}" != "${_rocm_env_variant}" ]]; then
67
+ echo "ERROR: VLLM_PRECOMPILED_WHEEL_VARIANT=${_rocm_variant} does not match detected environment ROCm variant ${_rocm_env_variant}" >&2
68
+ exit 1
69
+ fi
70
+ else
71
+ _rocm_variant="${_rocm_env_variant}"
72
+ fi
73
+ if [[ -z "${_rocm_variant}" ]]; then
74
+ echo "ERROR: Could not detect ROCm variant from the environment for commit ${merge_base_commit}" >&2
75
+ exit 1
76
+ fi
77
+ if [[ -z "${_available_variants}" ]] \
78
+ || [[ " ${_available_variants} " != *" ${_rocm_variant} "* ]]; then
79
+ echo "ERROR: Environment ROCm variant '${_rocm_variant}' is not published for commit ${merge_base_commit} (available:${_available_variants:-none})" >&2
80
+ exit 1
81
+ fi
82
+ meta_json_url="https://wheels.vllm.ai/rocm/${merge_base_commit}/${_rocm_variant}/vllm/metadata.json"
83
+ unset -v _rocm_env_variant _available_variants _rocm_variant
84
+else
85
+ meta_json_url="https://wheels.vllm.ai/${merge_base_commit}/vllm/metadata.json"
86
+fi
87
+unset -v _vllm_target_lower
88
+echo "INFO: will use metadata.json from ${meta_json_url}"
89
90
for i in {1..5}; do
91
echo "Checking metadata.json URL (attempt $i)..."