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 8, 17:45 UTC
[bugfix] [ROCm] Fix premature CUDA initialization in platform detection (#33941)
Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
kourosh hakhamaneshi committed
Feb 6, 2026 at 14:17 UTC
4a2d00eafd45d06a63936fcd4461b129382e783b
6 files changed
+133
-6
.buildkite/test-amd.yaml
+1
@@ -534,6 +534,7 @@ steps:
534
- tests/cuda
535
commands:
536
- pytest -v -s cuda/test_cuda_context.py
537
+ - pytest -v -s cuda/test_platform_no_cuda_init.py
538
539
- label: Samplers Test # 56min
540
timeout_in_minutes: 75
.buildkite/test_areas/cuda.yaml
+1
@@ -9,6 +9,7 @@ steps:
9
- tests/cuda
10
commands:
11
- pytest -v -s cuda/test_cuda_context.py
12
+ - pytest -v -s cuda/test_platform_no_cuda_init.py
13
14
- label: Cudagraph
15
timeout_in_minutes: 20
tests/cuda/scripts/check_device_count_respects_env.py
new
+23
@@ -0,0 +1,23 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""Check that device_count respects CUDA_VISIBLE_DEVICES after platform import."""
5
+
6
+import os
7
+import sys
8
+
9
+for key in ["CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"]:
10
+ os.environ.pop(key, None)
11
+
12
+import torch # noqa: E402
13
+
14
+from vllm.platforms import current_platform # noqa: F401, E402
15
+
16
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
17
+count = torch.cuda.device_count()
18
+
19
+if count == 0:
20
+ sys.exit(0) # Skip: no GPUs available
21
+
22
+assert count == 1, f"device_count()={count}, expected 1"
23
+print("OK")
tests/cuda/scripts/check_platform_no_cuda_init.py
new
+20
@@ -0,0 +1,20 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""Check that vllm.platforms import does not initialize CUDA."""
5
+
6
+import os
7
+
8
+for key in ["CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"]:
9
+ os.environ.pop(key, None)
10
+
11
+import torch # noqa: E402
12
+
13
+assert not torch.cuda.is_initialized(), "CUDA initialized before import"
14
+
15
+from vllm.platforms import current_platform # noqa: E402
16
+
17
+assert not torch.cuda.is_initialized(), (
18
+ f"CUDA was initialized during vllm.platforms import on {current_platform}"
19
+)
20
+print("OK")
tests/cuda/test_platform_no_cuda_init.py
new
+48
@@ -0,0 +1,48 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test that platform imports do not prematurely initialize CUDA.
5
+
6
+This is critical for Ray-based multi-GPU setups where workers need to
7
+set CUDA_VISIBLE_DEVICES after importing vLLM but before CUDA is initialized.
8
+If CUDA is initialized during import, device_count() gets locked and ignores
9
+subsequent env var changes.
10
+"""
11
+
12
+import subprocess
13
+import sys
14
+from pathlib import Path
15
+
16
+import pytest
17
+
18
+SCRIPTS_DIR = Path(__file__).parent / "scripts"
19
+
20
+
21
+def run_script(script_name: str) -> subprocess.CompletedProcess:
22
+ """Run a test script in a subprocess with clean CUDA state."""
23
+ script_path = SCRIPTS_DIR / script_name
24
+ return subprocess.run(
25
+ [sys.executable, str(script_path)],
26
+ capture_output=True,
27
+ text=True,
28
+ )
29
+
30
+
31
+def test_platform_import_does_not_init_cuda():
32
+ """Test that importing vllm.platforms does not initialize CUDA."""
33
+ result = run_script("check_platform_no_cuda_init.py")
34
+ if result.returncode != 0:
35
+ pytest.fail(f"Platform import initialized CUDA:\n{result.stderr}")
36
+
37
+
38
+def test_device_count_respects_env_after_platform_import():
39
+ """Test that device_count respects CUDA_VISIBLE_DEVICES after import."""
40
+ result = run_script("check_device_count_respects_env.py")
41
+ if result.returncode != 0:
42
+ pytest.fail(
43
+ f"device_count does not respect env var after import:\n{result.stderr}"
44
+ )
45
+
46
+
47
+if __name__ == "__main__":
48
+ pytest.main([__file__, "-v"])
vllm/platforms/rocm.py
+40
-6
@@ -87,33 +87,67 @@ def with_amdsmi_context(fn):
87
return wrapper
88
89
90
+@with_amdsmi_context
91
+def _query_gcn_arch_from_amdsmi() -> str:
92
+ """Query GCN arch from amdsmi. Raises if not available."""
93
+ handles = amdsmi_get_processor_handles()
94
+ if handles:
95
+ asic_info = amdsmi_get_gpu_asic_info(handles[0])
96
+ # Use target_graphics_version which contains the gfx name
97
+ # e.g., 'gfx942' for MI300X/MI325X
98
+ target_gfx = asic_info.get("target_graphics_version", "")
99
+ if target_gfx:
100
+ return target_gfx
101
+ raise RuntimeError("amdsmi did not return valid GCN arch")
102
+
103
+
104
+@cache
105
+def _get_gcn_arch_via_amdsmi() -> str:
106
+ """
107
+ Get the GCN architecture name using amdsmi instead of torch.cuda.
108
+ This avoids initializing CUDA, which is important for Ray workers
109
+ that need to set CUDA_VISIBLE_DEVICES after importing vLLM.
110
+ """
111
+ try:
112
+ return _query_gcn_arch_from_amdsmi()
113
+ except Exception as e:
114
+ logger.debug("Failed to get GCN arch via amdsmi: %s", e)
115
+ logger.warning_once(
116
+ "Failed to get GCN arch via amdsmi, falling back to torch.cuda. "
117
+ "This will initialize CUDA and may cause "
118
+ "issues if CUDA_VISIBLE_DEVICES is not set yet."
119
+ )
120
+ # Ultimate fallback: use torch.cuda (will initialize CUDA)
121
+ return torch.cuda.get_device_properties("cuda").gcnArchName
122
+
123
+
124
@cache
125
def on_gfx1x() -> bool:
92
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
126
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
127
return any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
128
129
130
@cache
131
def on_mi3xx() -> bool:
98
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
132
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
133
return any(arch in GPU_ARCH for arch in ["gfx942", "gfx950"])
134
135
136
@cache
137
def on_gfx9() -> bool:
104
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
138
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
139
return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
140
141
142
@cache
143
def on_gfx942() -> bool:
110
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
144
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
145
return any(arch in GPU_ARCH for arch in ["gfx942"])
146
147
148
@cache
149
def on_gfx950() -> bool:
116
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
150
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
151
return any(arch in GPU_ARCH for arch in ["gfx950"])
152
153
@@ -129,7 +163,7 @@ def use_rocm_custom_paged_attention(
163
alibi_slopes: torch.Tensor | None = None,
164
sinks: torch.Tensor | None = None,
165
) -> bool:
132
- GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
166
+ GPU_ARCH = _get_gcn_arch_via_amdsmi()
167
ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
168
ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
169