Fix: keep PYTORCH_NVML_BASED_CUDA_CHECK=1 on MIG

Setting PYTORCH_NVML_BASED_CUDA_CHECK=0 on MIG causes PyTorch to
use cudaGetDeviceCount() instead of NVML for device checks, which
initializes CUDA at import time. This breaks EngineCore's forked
subprocesses with "Cannot re-initialize CUDA in forked subprocess".

Keep NVML-based checks enabled (=1) on all platforms. Only set
expandable_segments:False on MIG.

Signed-off-by: khluu <khluu000@gmail.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/test_areas/models_basic.yaml

@@ -7,6 +7,7 @@ steps:
   timeout_in_minutes: 45
   device: h200_18gb
   torch_nightly: true
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/test_initialization.py

.buildkite/test_areas/models_language.yaml

@@ -6,6 +6,7 @@ steps:
   key: language-models-tests-standard
   timeout_in_minutes: 25
   device: h200_18gb
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/language

vllm/env_override.py

@@ -100,6 +100,18 @@ logger = init_logger(__name__)
 # it avoids unintentional cuda initialization from torch.cuda.is_available()
 os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
 
+# On MIG partitions, NVML access is restricted — PyTorch's
+# CUDACachingAllocator hits an NVML assertion when it calls NVML
+# internally during pool growth. Disable expandable_segments to
+# reduce (but not eliminate) NVML-dependent code paths.
+# Note: PYTORCH_NVML_BASED_CUDA_CHECK must stay "1" even on MIG,
+# because setting it to "0" causes early CUDA initialization that
+# breaks forked subprocesses.
+_nvidia_visible = os.environ.get("NVIDIA_VISIBLE_DEVICES", "")
+if _nvidia_visible.startswith("MIG-"):
+    os.environ.setdefault(
+        "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False")
+
 # see https://github.com/vllm-project/vllm/issues/10480 and
 # https://github.com/vllm-project/vllm/issues/10619.
 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"

vllm/platforms/cuda.py

@@ -167,6 +167,16 @@ class CudaPlatformBase(Platform):
         "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
     ]
 
+    @classmethod
+    def is_mig(cls) -> bool:
+        """Check if running on a MIG (Multi-Instance GPU) partition.
+
+        MIG partitions have restricted NVML access which can cause
+        PyTorch's CUDACachingAllocator to fail with NVML assertions.
+        """
+        nvidia_visible = os.environ.get("NVIDIA_VISIBLE_DEVICES", "")
+        return nvidia_visible.startswith("MIG-")
+
     @property
     def supported_dtypes(self) -> list[torch.dtype]:
         if self.has_device_capability(80):