[Bugfix] Add MigCudaPlatform to avoid CUDA init before fork

NonNvmlCudaPlatform uses torch.cuda for get_device_capability(),
which initializes the CUDA context. When this runs at import time
(e.g. via requires_fp8 in test utils), it prevents fork-based tests
from working ("Cannot re-initialize CUDA in forked subprocess").

MigCudaPlatform inherits from NvmlCudaPlatform (NVML-based capability
queries don't init CUDA) and only overrides the memory methods that
fail on MIG partitions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Signed-off-by: khluu <khluu000@gmail.com>

tests/utils.py

@@ -383,9 +383,14 @@ class RemoteVLLMServer:
                     total_used = 0
                     device_count = current_platform.device_count()
                     for i in range(device_count):
-                        handle = nvmlDeviceGetHandleByIndex(i)
-                        mem_info = nvmlDeviceGetMemoryInfo(handle)
-                        total_used += mem_info.used
+                        try:
+                            handle = nvmlDeviceGetHandleByIndex(i)
+                            mem_info = nvmlDeviceGetMemoryInfo(handle)
+                            total_used += mem_info.used
+                        except Exception:
+                            # NVML fails on MIG — fall back to torch.cuda
+                            free, total = torch.cuda.mem_get_info(i)
+                            total_used += total - free
                     return total_used
         except Exception as e:
             print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}")
@@ -1202,7 +1207,7 @@ def wait_for_gpu_memory_to_clear(
 ) -> None:
     assert threshold_bytes is not None or threshold_ratio is not None
     # Use nvml instead of pytorch to reduce measurement error from torch cuda
-    # context.
+    # context. Falls back to torch.cuda on MIG where NVML lacks permissions.
     devices = get_physical_device_indices(devices)
     start_time = time.time()
     while True:
@@ -1215,10 +1220,17 @@ def wait_for_gpu_memory_to_clear(
                 gb_used = mem_info["vram_used"] / 2**10
                 gb_total = mem_info["vram_total"] / 2**10
             else:
-                dev_handle = nvmlDeviceGetHandleByIndex(device)
-                mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
-                gb_used = mem_info.used / 2**30
-                gb_total = mem_info.total / 2**30
+                try:
+                    dev_handle = nvmlDeviceGetHandleByIndex(device)
+                    mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
+                    gb_used = mem_info.used / 2**30
+                    gb_total = mem_info.total / 2**30
+                except Exception:
+                    # NVML fails on MIG with NoPermission — fall back to
+                    # torch.cuda which correctly reports MIG partition memory.
+                    free, total = torch.cuda.mem_get_info(device)
+                    gb_total = total / 2**30
+                    gb_used = (total - free) / 2**30
             output_raw[device] = (gb_used, gb_total)
             output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
 

vllm/platforms/cuda.py

@@ -844,6 +844,37 @@ class NonNvmlCudaPlatform(CudaPlatformBase):
         return None
 
 
+class MigCudaPlatform(NvmlCudaPlatform):
+    """Platform for MIG (Multi-Instance GPU) devices.
+
+    Inherits from NvmlCudaPlatform so that device queries like
+    get_device_capability() and get_device_name() use NVML instead of
+    torch.cuda, avoiding premature CUDA context initialization (which
+    would break fork-based test frameworks).
+
+    Overrides only the methods that fail on MIG partitions due to NVML
+    NoPermission errors (memory queries report full-GPU memory instead
+    of partition memory).
+    """
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        # NVML returns the *full* GPU's memory on MIG partitions, not the
+        # partition's memory. Use torch.cuda which reports correctly.
+        device_props = torch.cuda.get_device_properties(device_id)
+        return device_props.total_memory
+
+    @classmethod
+    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
+        # P2P / NVLink queries are not meaningful on MIG partitions.
+        return False
+
+    @classmethod
+    def log_warnings(cls):
+        # Skip NVML-based heterogeneous-GPU warnings on MIG partitions.
+        pass
+
+
 # Autodetect either NVML-enabled or non-NVML platform
 # based on whether NVML is available.
 nvml_available = False
@@ -858,6 +889,83 @@ finally:
     if nvml_available:
         pynvml.nvmlShutdown()
 
-CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
+
+def _is_mig_device() -> bool:
+    """Detect if the current CUDA device is a MIG (Multi-Instance GPU) partition.
+
+    On MIG partitions, most NVML device-level queries fail with
+    NVMLError_NoPermission. This causes PyTorch's CUDACachingAllocator to crash
+    with an internal assert (via expandable_segments) before the real OOM error
+    can surface. It also causes NVML memory queries to report the full GPU's
+    memory instead of the MIG partition's memory.
+
+    Detection methods:
+    1. CUDA_VISIBLE_DEVICES containing MIG UUIDs (MIG-GPU-...)
+    2. Probing NVML device queries for permission errors
+    """
+    # Method 1: Check CUDA_VISIBLE_DEVICES for MIG UUID format
+    cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+    if cuda_visible and any(
+        dev.strip().startswith("MIG-") for dev in cuda_visible.split(",")
+    ):
+        logger.info(
+            "Detected MIG device from CUDA_VISIBLE_DEVICES=%s, "
+            "using MigCudaPlatform to avoid NVML memory errors.",
+            cuda_visible,
+        )
+        return True
+
+    # Method 2: Probe NVML for permission errors (catches MIG even when
+    # CUDA_VISIBLE_DEVICES uses numeric indices mapped by the container runtime)
+    if nvml_available:
+        try:
+            pynvml.nvmlInit()
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+                pynvml.nvmlDeviceGetMemoryInfo(handle)
+            except pynvml.NVMLError as e:
+                if "NoPermission" in str(
+                    type(e).__name__
+                ) or "Insufficient Permissions" in str(e):
+                    logger.info(
+                        "Detected MIG device via NVML permission error: %s, "
+                        "using MigCudaPlatform.",
+                        e,
+                    )
+                    return True
+            finally:
+                pynvml.nvmlShutdown()
+        except Exception:
+            pass
+
+    return False
+
+
+_mig_device = _is_mig_device()
+
+CudaPlatform: type[NvmlCudaPlatform] | type[NonNvmlCudaPlatform] | type[MigCudaPlatform]
+if _mig_device:
+    CudaPlatform = MigCudaPlatform
+
+    # Disable expandable_segments on MIG devices. PyTorch's
+    # CUDACachingAllocator internally calls NVML (for the
+    # expandable_segments feature) which fails on MIG with
+    # "NVML_SUCCESS == r INTERNAL ASSERT FAILED", crashing before
+    # a real torch.cuda.OutOfMemoryError can surface.
+    _alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+    if "expandable_segments" not in _alloc_conf:
+        new_val = (
+            f"{_alloc_conf},expandable_segments:False"
+            if _alloc_conf
+            else "expandable_segments:False"
+        )
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = new_val
+        logger.info(
+            "MIG device detected: set PYTORCH_CUDA_ALLOC_CONF=%s "
+            "to prevent NVML assert in CUDACachingAllocator.",
+            new_val,
+        )
+else:
+    CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
 
 CudaPlatform.log_warnings()