[ROCm][CI] Query total device memory via amdsmi to avoid HIP init (#46141)
Signed-off-by: stefankoncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Jun 22, 2026 at 22:12 UTC
2b4a7491ecffc362f1f080c3ac41dfe10018c39b
1 file changed
+26
-2
vllm/platforms/rocm.py
+26
-2
@@ -27,8 +27,10 @@ logger = init_logger(__name__)
27
try:
28
from amdsmi import (
29
AmdSmiException,
30
+ AmdSmiMemoryType,
31
amdsmi_get_gpu_asic_info,
32
amdsmi_get_gpu_device_uuid,
33
+ amdsmi_get_gpu_memory_total,
34
amdsmi_get_processor_handles,
35
amdsmi_init,
36
amdsmi_shut_down,
@@ -167,6 +169,14 @@ def _query_gcn_arch_from_amdsmi() -> str:
169
raise RuntimeError("amdsmi did not return valid GCN arch")
170
171
172
+@with_amdsmi_context
173
+def _query_total_memory_from_amdsmi(physical_device_id: int) -> int:
174
+ """Query total VRAM (bytes) from amdsmi. Raises if not available."""
175
+ handles = amdsmi_get_processor_handles()
176
+ handle = handles[physical_device_id]
177
+ return amdsmi_get_gpu_memory_total(handle, AmdSmiMemoryType.VRAM)
178
+
179
+
180
def _get_gcn_arch() -> str:
181
"""
182
Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda.
@@ -726,8 +736,22 @@ class RocmPlatform(Platform):
736
737
@classmethod
738
def get_device_total_memory(cls, device_id: int = 0) -> int:
729
- device_props = torch.cuda.get_device_properties(device_id)
730
- return device_props.total_memory
739
+ # Query total VRAM via amdsmi so we don't initialize a HIP context in
740
+ # the calling process. torch.cuda.get_device_properties() creates a
741
+ # HIP context, which makes vLLM fall back from `fork` to `spawn` for
742
+ # worker processes. Keeping this query context-free preserves `fork`
743
+ # where it is otherwise valid (e.g. out-of-tree models registered in
744
+ # the parent process).
745
+ try:
746
+ physical_device_id = cls.device_id_to_physical_device_id(device_id)
747
+ return _query_total_memory_from_amdsmi(physical_device_id)
748
+ except Exception as e:
749
+ logger.debug("Failed to get total memory via amdsmi: %s", e)
750
+ logger.warning_once(
751
+ "Failed to get total memory via amdsmi, falling back to "
752
+ "torch.cuda. This will initialize CUDA."
753
+ )
754
+ return torch.cuda.get_device_properties(device_id).total_memory
755
756
@classmethod
757
def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: