diff --git a/tests/cuda/test_cuda_context.py b/tests/cuda/test_cuda_context.py index 6336f2112c6..16d2f16c2d8 100644 --- a/tests/cuda/test_cuda_context.py +++ b/tests/cuda/test_cuda_context.py @@ -77,5 +77,43 @@ class TestSetCudaContext: current_platform.set_device(torch.device("cpu")) +def test_get_device_capability_uses_visible_device_ordinal(monkeypatch): + import vllm.platforms.interface as platform_interface + from vllm.platforms.cuda import NvmlCudaPlatform, pynvml + + seen_indices: list[int] = [] + + def record_handle(index: int) -> str: + seen_indices.append(index) + return f"handle-{index}" + + monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [1]) + monkeypatch.setenv(NvmlCudaPlatform.device_control_env_var, "0,1") + monkeypatch.setattr( + NvmlCudaPlatform, + "device_control_id_to_physical_device_id", + classmethod(lambda _cls, device_id: int(device_id)), + ) + monkeypatch.setattr(pynvml, "nvmlInit", lambda: None) + monkeypatch.setattr(pynvml, "nvmlShutdown", lambda: None) + monkeypatch.setattr( + pynvml, + "nvmlDeviceGetHandleByIndex", + record_handle, + ) + monkeypatch.setattr( + pynvml, + "nvmlDeviceGetCudaComputeCapability", + lambda _handle: (9, 0), + ) + NvmlCudaPlatform.get_device_capability.cache_clear() + + capability = NvmlCudaPlatform.get_device_capability(device_id=1) + + assert capability is not None + assert capability.to_int() == 90 + assert seen_indices == [1] + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 605650525af..9eac95e0324 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -727,7 +727,7 @@ class NvmlCudaPlatform(CudaPlatformBase): @with_nvml_context def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: try: - physical_device_id = cls.device_id_to_physical_device_id(device_id) + physical_device_id = cls.visible_device_id_to_physical_device_id(device_id) handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id) major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) return DeviceCapability(major=major, minor=minor) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index a81c34d7a50..7cf009c8071 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -271,6 +271,16 @@ class Platform: f"{cls.device_name}." ) from e + # GPU device IDs can refer to three distinct namespaces: + # - logical: vLLM-local IDs such as local ranks. These index + # assigned_physical_gpu_ids when it is set. + # - visible: torch/CUDA ordinals in the current process after applying + # the device-control env var, e.g. CUDA_VISIBLE_DEVICES. + # - physical: global GPU IDs used by topology and management APIs such as + # NVML, which are not remapped by CUDA_VISIBLE_DEVICES. + # Keep conversions explicit. In particular, torch device indices are + # visible IDs, not vLLM logical IDs. + @classmethod def device_id_to_physical_device_id(cls, device_id: int): """Map a vLLM-local logical device ID to a physical device ID. @@ -411,7 +421,12 @@ class Platform: cls, device_id: int = 0, ) -> DeviceCapability | None: - """Stateless version of [torch.cuda.get_device_capability][].""" + """Stateless version of [torch.cuda.get_device_capability][]. + + Args: + device_id: Device index in the visible device namespace, matching + the argument accepted by torch.cuda. + """ return None @classmethod