diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9662037b01f..06953d504b6 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -27,8 +27,10 @@ logger = init_logger(__name__) try: from amdsmi import ( AmdSmiException, + AmdSmiMemoryType, amdsmi_get_gpu_asic_info, amdsmi_get_gpu_device_uuid, + amdsmi_get_gpu_memory_total, amdsmi_get_processor_handles, amdsmi_init, amdsmi_shut_down, @@ -167,6 +169,14 @@ def _query_gcn_arch_from_amdsmi() -> str: raise RuntimeError("amdsmi did not return valid GCN arch") +@with_amdsmi_context +def _query_total_memory_from_amdsmi(physical_device_id: int) -> int: + """Query total VRAM (bytes) from amdsmi. Raises if not available.""" + handles = amdsmi_get_processor_handles() + handle = handles[physical_device_id] + return amdsmi_get_gpu_memory_total(handle, AmdSmiMemoryType.VRAM) + + def _get_gcn_arch() -> str: """ Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda. @@ -726,8 +736,22 @@ class RocmPlatform(Platform): @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: - device_props = torch.cuda.get_device_properties(device_id) - return device_props.total_memory + # Query total VRAM via amdsmi so we don't initialize a HIP context in + # the calling process. torch.cuda.get_device_properties() creates a + # HIP context, which makes vLLM fall back from `fork` to `spawn` for + # worker processes. Keeping this query context-free preserves `fork` + # where it is otherwise valid (e.g. out-of-tree models registered in + # the parent process). + try: + physical_device_id = cls.device_id_to_physical_device_id(device_id) + return _query_total_memory_from_amdsmi(physical_device_id) + except Exception as e: + logger.debug("Failed to get total memory via amdsmi: %s", e) + logger.warning_once( + "Failed to get total memory via amdsmi, falling back to " + "torch.cuda. This will initialize CUDA." + ) + return torch.cuda.get_device_properties(device_id).total_memory @classmethod def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None: