A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 14:00 UTC
[ROCm] Derive device capability from GCN arch string without CUDA init (#35069)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 27, 2026 at 23:55 UTC
94029ffaf02f0b73e296e11cab721c23fd5a5f97
2 files changed
+120
-7
vllm/platforms/rocm.py
+104
-7
@@ -6,6 +6,7 @@ from datetime import timedelta
6
from functools import cache, lru_cache, wraps
7
from typing import TYPE_CHECKING
8
9
+import regex as re
10
import torch
11
from torch.distributed import PrefixStore, ProcessGroup
12
from torch.distributed.distributed_c10d import is_nccl_available
@@ -64,13 +65,29 @@ _ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {
65
"0x744c": "AMD_Radeon_RX7900XTX",
66
}
67
67
-# Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES`
68
-if "HIP_VISIBLE_DEVICES" in os.environ:
69
- val = os.environ["HIP_VISIBLE_DEVICES"]
70
- if cuda_val := os.environ.get("CUDA_VISIBLE_DEVICES", None):
71
- assert val == cuda_val
72
- else:
73
- os.environ["CUDA_VISIBLE_DEVICES"] = val
68
+
69
+def _sync_hip_cuda_env_vars():
70
+ """Ensure HIP_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES are consistent.
71
+ Treats empty string as unset. Raises on genuine conflicts."""
72
+ hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None
73
+ cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None
74
+
75
+ if hip_val is not None and cuda_val is not None:
76
+ if hip_val != cuda_val:
77
+ raise ValueError(
78
+ f"Inconsistent GPU visibility env vars: "
79
+ f"HIP_VISIBLE_DEVICES='{hip_val}' vs "
80
+ f"CUDA_VISIBLE_DEVICES='{cuda_val}'. "
81
+ f"Please set only one, or ensure they match."
82
+ )
83
+ elif hip_val is not None:
84
+ os.environ["CUDA_VISIBLE_DEVICES"] = hip_val
85
+ elif cuda_val is not None:
86
+ os.environ["HIP_VISIBLE_DEVICES"] = cuda_val
87
+
88
+
89
+# Sync at import time - catches misconfigurations from process start.
90
+_sync_hip_cuda_env_vars()
91
92
# AMDSMI utils
93
# Note that NVML is not affected by `{CUDA/HIP}_VISIBLE_DEVICES`,
@@ -134,6 +151,77 @@ _ON_GFX942 = "gfx942" in _GCN_ARCH
151
_ON_GFX950 = "gfx950" in _GCN_ARCH
152
153
154
+def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:
155
+ """
156
+ Parse (major, minor) from a GCN arch string, mirroring how
157
+ HIP derives hipDeviceProp_t.major / .minor.
158
+
159
+ Format: gfx<MAJOR><MINOR><STEPPING>
160
+ - 1-digit major (gfx9xx): "gfx" + M + m + stepping
161
+ - 2-digit major (gfx1xxx): "gfx" + MM + m + stepping
162
+
163
+ Examples:
164
+ gfx90a -> (9, 0) gfx942 -> (9, 4) gfx950 -> (9, 5)
165
+ gfx1100 -> (11, 0) gfx1101 -> (11, 0) gfx1200 -> (12, 0)
166
+
167
+ Returns None only when the string is not gfx-prefixed at all
168
+ (i.e. not a ROCm arch string). Raises on any string that looks
169
+ like a GCN arch but does not match a known layout.
170
+ """
171
+ m = re.match(r"gfx(\d+)", gcn_arch)
172
+ if not m:
173
+ # Not a gfx string at all — caller should fall back to torch.cuda
174
+ return None
175
+
176
+ digits = m.group(1)
177
+ n = len(digits)
178
+
179
+ if n < 2:
180
+ raise ValueError(
181
+ f"GCN arch '{gcn_arch}' has too few digits ({n}) after 'gfx' "
182
+ f"to derive a (major, minor) capability. "
183
+ f"Please file a vLLM issue with your GPU model."
184
+ )
185
+
186
+ if n in (2, 3):
187
+ # 1-digit major: gfx9 family
188
+ # len 2: major + minor (e.g. gfx90 from gfx90a)
189
+ # len 3: major + minor + step (e.g. gfx942)
190
+ major = int(digits[0])
191
+ minor = int(digits[1])
192
+ elif n == 4:
193
+ # 2-digit major: gfx10xx, gfx11xx, gfx12xx
194
+ # major(2) + minor(1) + stepping(1)
195
+ major = int(digits[:2])
196
+ minor = int(digits[2])
197
+ elif n >= 5:
198
+ raise ValueError(
199
+ f"GCN arch '{gcn_arch}' has {n} digits after 'gfx', which "
200
+ f"exceeds the known 4-digit layout (MMms). Cannot determine "
201
+ f"major/minor split unambiguously. "
202
+ f"Please file a vLLM issue with your GPU model."
203
+ )
204
+
205
+ if major < 9:
206
+ raise ValueError(
207
+ f"Parsed unknown ROCm architecture from GCN arch '{gcn_arch}': "
208
+ f"major={major}, minor={minor}. "
209
+ f"Major version < 9 is not expected for any supported AMD GPU. "
210
+ f"Please file a vLLM issue with your GPU model."
211
+ )
212
+
213
+ if major > 12:
214
+ raise ValueError(
215
+ f"Parsed unknown ROCm architecture from GCN arch '{gcn_arch}': "
216
+ f"major={major}, minor={minor}. "
217
+ f"Major version > 12 is beyond currently known AMD generations. "
218
+ f"Please file a vLLM issue with your GPU model so support "
219
+ f"can be added."
220
+ )
221
+
222
+ return (major, minor)
223
+
224
+
225
def on_gfx1x() -> bool:
226
return _ON_GFX1X
227
@@ -444,6 +532,15 @@ class RocmPlatform(Platform):
532
@classmethod
533
@lru_cache(maxsize=8)
534
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
535
+ cap = _capability_from_gcn_arch(_GCN_ARCH)
536
+ if cap is not None:
537
+ return DeviceCapability(major=cap[0], minor=cap[1])
538
+
539
+ logger.warning_once(
540
+ "Could not derive device capability from GCN arch '%s', "
541
+ "falling back to torch.cuda (this will initialize CUDA).",
542
+ _GCN_ARCH,
543
+ )
544
major, minor = torch.cuda.get_device_capability(device_id)
545
return DeviceCapability(major=major, minor=minor)
546
vllm/utils/system_utils.py
+16
@@ -16,6 +16,7 @@ import psutil
16
17
import vllm.envs as envs
18
from vllm.logger import init_logger
19
+from vllm.platforms import current_platform
20
from vllm.platforms.interface import in_wsl
21
from vllm.ray.lazy_utils import is_in_ray_actor
22
@@ -111,6 +112,17 @@ def unique_filepath(fn: Callable[[int], Path]) -> Path:
112
# Process management utilities
113
114
115
+def _sync_visible_devices_env_vars():
116
+ """Sync HIP/CUDA visibility env vars before spawning (ROCm only)."""
117
+
118
+ if not current_platform.is_rocm():
119
+ return
120
+
121
+ from vllm.platforms.rocm import _sync_hip_cuda_env_vars
122
+
123
+ _sync_hip_cuda_env_vars()
124
+
125
+
126
def _maybe_force_spawn():
127
"""Check if we need to force the use of the `spawn` multiprocessing start
128
method.
@@ -156,6 +168,10 @@ def get_mp_context():
168
VLLM_WORKER_MULTIPROC_METHOD.
169
"""
170
_maybe_force_spawn()
171
+ # (ROCm): Sync GPU visibility env vars so spawned children inherit
172
+ # consistent values. Must run after _maybe_force_spawn and regardless
173
+ # of whether spawn was already set.
174
+ _sync_visible_devices_env_vars()
175
mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
176
return multiprocessing.get_context(mp_method)
177