From 154aa68e32745aa9be58aec044e0069663447007 Mon Sep 17 00:00:00 2001 From: khluu Date: Wed, 13 May 2026 04:21:55 -0700 Subject: [PATCH] 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 Co-Authored-By: Claude Opus 4.6 (1M context) --- .buildkite/test_areas/kernels.yaml | 1 - .buildkite/test_areas/lm_eval.yaml | 1 - .buildkite/test_areas/misc.yaml | 1 - vllm/env_override.py | 11 +++++++---- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 6ac7ae0ae78..fee7728938f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.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 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f001b5cefb5..a9e76b669f0 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.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 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index cf529925539..ca9d3fcef82 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.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 diff --git a/vllm/env_override.py b/vllm/env_override.py index a14330a8ae9..6c57f06bbff 100644 --- a/vllm/env_override.py +++ b/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.