Fix NVML capability lookup for visible devices (#47892)

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

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"])

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)

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