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/kernels.yaml

@@ -92,7 +92,6 @@ steps:
 - label: Kernels KDA Test
   timeout_in_minutes: 20
   device: h200_18gb
-  device: h200_18gb
   source_file_dependencies:
   - vllm/model_executor/layers/fla/ops/kda.py
   - vllm/model_executor/layers/fla/ops/chunk_delta_h.py

.buildkite/test_areas/lm_eval.yaml

@@ -115,7 +115,6 @@ steps:
   key: lm-eval-turboquant-kv-cache
   timeout_in_minutes: 75
   device: h200_18gb
-  device: h200_18gb
   source_file_dependencies:
   - vllm/model_executor/layers/quantization/turboquant/
   - vllm/v1/attention/backends/turboquant_attn.py

.buildkite/test_areas/misc.yaml

@@ -215,7 +215,6 @@ steps:
   optional: true
   timeout_in_minutes: 20
   device: h200_18gb
-  device: h200_18gb
   source_file_dependencies:
   - tests/standalone_tests/python_only_compile.sh
   - setup.py

vllm/env_override.py

@@ -98,16 +98,19 @@ logger = init_logger(__name__)
 
 # see https://github.com/vllm-project/vllm/pull/15951
 # 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. Use CUDA runtime checks instead.
+# 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_NVML_BASED_CUDA_CHECK", "0")
     os.environ.setdefault(
         "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:False")
-else:
-    os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
 
 # see https://github.com/vllm-project/vllm/issues/10480 and
 # https://github.com/vllm-project/vllm/issues/10619.