Remove broken allocator pre-warming (initializes CUDA before fork)

The pre-warming called torch.cuda.mem_get_info() in the parent process before the EngineCore subprocess forks, which initialized CUDA. This caused "Cannot re-initialize CUDA in forked subprocess" in the child process, breaking tests that previously passed. The pre-warming also wouldn't have helped — PyTorch's allocator pool is per-process, so the child's pool starts empty regardless. Signed-off-by: khluu <khluu000@gmail.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
[CI] Migrate all remaining gpu_1_queue jobs to h200_18gb MIG
2026-05-13 02:26:43 -07:00 · 2026-05-12 18:50:08 -07:00 · 2026-05-12 15:57:19 -07:00
17 changed files with 68 additions and 1 deletions
@@ -20,6 +20,7 @@ steps:
 - label: Cudagraph
  key: cudagraph
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - tests/v1/cudagraph
  - vllm/v1/cudagraph_dispatcher.py
@@ -30,6 +30,7 @@ steps:
 - label: Engine (1 GPU)
  key: engine-1-gpu
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/engine/
    - tests/v1/engine/
@@ -50,6 +51,7 @@ steps:
 - label: e2e Core (1 GPU)
  key: e2e-core-1-gpu
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/
    - tests/v1/e2e/general/
@@ -6,6 +6,7 @@ steps:
  key: entrypoints-unit-tests
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/entrypoints
  - tests/entrypoints/
@@ -17,6 +18,7 @@ steps:
  key: entrypoints-integration-llm
  timeout_in_minutes: 40
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/llm
@@ -36,6 +38,7 @@ steps:
  key: entrypoints-integration-api-server-openai-part-1
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai
@@ -48,6 +51,7 @@ steps:
  key: entrypoints-integration-api-server-openai-part-2
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai
@@ -73,6 +77,7 @@ steps:
  key: entrypoints-integration-api-server-2
  timeout_in_minutes: 130
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/rpc
@@ -88,6 +93,7 @@ steps:
  key: entrypoints-integration-speech_to_text
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/speech_to_text
@@ -99,6 +105,7 @@ steps:
  key: entrypoints-integration-pooling
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/pooling
@@ -110,6 +117,7 @@ steps:
  key: entrypoints-integration-responses-api
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/entrypoints/openai/responses
@@ -17,6 +17,7 @@ steps:
 - label: Kernels Core Operation Test
  key: kernels-core-operation-test
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - tests/kernels/core
@@ -41,6 +42,7 @@ steps:
 - label: Kernels Attention Test %N
  key: kernels-attention-test
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
  - csrc/attention/
  - vllm/v1/attention
@@ -55,6 +57,7 @@ steps:
 - label: Kernels Quantization Test %N
  key: kernels-quantization-test
  timeout_in_minutes: 90
+  device: h200_18gb
  source_file_dependencies:
  - csrc/quantization/
  - vllm/model_executor/layers/quantization
@@ -66,6 +69,7 @@ steps:
 - label: Kernels MoE Test %N
  key: kernels-moe-test
  timeout_in_minutes: 25
+  device: h200_18gb
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
  - csrc/moe/
@@ -82,6 +86,7 @@ steps:
 - label: Kernels Mamba Test
  key: kernels-mamba-test
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - csrc/mamba/
  - tests/kernels/mamba
@@ -91,6 +96,7 @@ steps:

 - label: Kernels KDA Test
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/fla/ops/kda.py
  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
@@ -5,6 +5,7 @@ steps:
 - label: LM Eval Small Models
  key: lm-eval-small-models
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -114,6 +115,7 @@ steps:
 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization/turboquant/
  - vllm/v1/attention/backends/turboquant_attn.py
@@ -152,6 +154,7 @@ steps:

 - label: MRCR Eval Small Models
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
  - tests/evals/mrcr/
  commands:
@@ -5,6 +5,7 @@ steps:
 - label: LoRA %N
  key: lora
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
  - vllm/lora
  - tests/lora
@@ -5,6 +5,7 @@ steps:
 - label: V1 Spec Decode
  key: v1-spec-decode
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -58,6 +59,7 @@ steps:
 - label: V1 Core + KV + Metrics
  key: v1-core-kv-metrics
  timeout_in_minutes: 30
+  device: h200_18gb
  source_file_dependencies:
    - vllm/config/
    - vllm/distributed/
@@ -153,6 +155,7 @@ steps:
  key: examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/entrypoints
  - vllm/multimodal
@@ -212,6 +215,7 @@ steps:
  depends_on: ~
  optional: true
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
  - tests/standalone_tests/python_only_compile.sh
  - setup.py
@@ -221,6 +225,7 @@ steps:
 - label: Async Engine, Inputs, Utils, Worker
  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 50
+  device: h200_18gb
  source_file_dependencies:
  - vllm/assets/
  - vllm/config/
@@ -5,6 +5,7 @@ steps:
 - label: Model Executor
  key: model-executor
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
@@ -5,6 +5,7 @@ steps:
 - label: Model Runner V2 Core Tests
  key: model-runner-v2-core-tests
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
@@ -29,6 +30,7 @@ steps:
  key: model-runner-v2-examples
  timeout_in_minutes: 45
  working_dir: "/vllm-workspace/examples"
+  device: h200_18gb
  source_file_dependencies:
    - vllm/v1/worker/gpu/
    - vllm/v1/core/sched/
@@ -102,6 +104,7 @@ steps:
  key: model-runner-v2-spec-decode
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/tests"
+  device: h200_18gb
  source_file_dependencies:
  - vllm/v1/worker/gpu/
  - vllm/v1/worker/gpu_worker.py
@@ -6,6 +6,7 @@ steps:
  key: basic-models-tests-initialization
  timeout_in_minutes: 45
  torch_nightly: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/test_initialization.py
@@ -19,6 +20,7 @@ steps:
 - label: Basic Models Tests (Extra Initialization) %N
  key: basic-models-tests-extra-initialization
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/test_initialization.py
@@ -35,6 +37,7 @@ steps:
 - label: Basic Models Tests (Other)
  key: basic-models-tests-other
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/test_terratorch.py
@@ -61,6 +64,7 @@ steps:
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
+  device: h200_18gb
  commands:
    - pip install --upgrade git+https://github.com/huggingface/transformers
    - pytest -v -s tests/models/test_initialization.py
@@ -77,6 +81,7 @@ steps:
  working_dir: "/vllm-workspace/"
  optional: true
  soft_fail: true
+  device: h200_18gb
  commands:
    - pip install transformers==4.57.5
    - pytest -v -s tests/models/test_initialization.py
@@ -5,6 +5,7 @@ steps:
 - label: Language Models Tests (Standard)
  key: language-models-tests-standard
  timeout_in_minutes: 25
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language
@@ -18,6 +19,7 @@ steps:
 - label: Language Models Tests (Extra Standard) %N
  key: language-models-tests-extra-standard
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/models/
  - tests/models/language/pooling/test_embedding.py
@@ -35,6 +37,7 @@ steps:
 - label: Language Models Tests (Hybrid) %N
  key: language-models-tests-hybrid
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/language/generation
@@ -60,6 +63,7 @@ steps:
 - label: Language Models Test (Extended Generation) # 80min
  key: language-models-test-extended-generation
  timeout_in_minutes: 110
+  device: h200_18gb
  optional: true
  source_file_dependencies:
  - vllm/
@@ -85,6 +89,7 @@ steps:
 - label: Language Models Test (Extended Pooling)  # 36min
  key: language-models-test-extended-pooling
  timeout_in_minutes: 50
+  device: h200_18gb
  optional: true
  source_file_dependencies:
  - vllm/
@@ -40,6 +40,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
  key: multi-modal-models-standard-3-llava-qwen2-vl
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -56,6 +57,7 @@ steps:
 - label: "Multi-Modal Models (Standard) 4: other + whisper"
  key: multi-modal-models-standard-4-other-whisper
  timeout_in_minutes: 45
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal
@@ -94,6 +96,7 @@ steps:
 - label: Multi-Modal Accuracy Eval (Small Models) # 50min
  key: multi-modal-accuracy-eval-small-models
  timeout_in_minutes: 70
+  device: h200_18gb
  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
  source_file_dependencies:
  - vllm/multimodal/
@@ -105,6 +108,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 1)
  key: multi-modal-models-extended-generation-1
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -122,6 +126,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 2)
  key: multi-modal-models-extended-generation-2
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -132,6 +137,7 @@ steps:
 - label: Multi-Modal Models (Extended Generation 3)
  key: multi-modal-models-extended-generation-3
  optional: true
+  device: h200_18gb
  source_file_dependencies:
  - vllm/
  - tests/models/multimodal/generation
@@ -5,6 +5,7 @@ steps:
 - label: PyTorch Compilation Unit Tests
  key: pytorch-compilation-unit-tests
  timeout_in_minutes: 10
+  device: h200_18gb
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -78,6 +79,7 @@ steps:
 - label: PyTorch Compilation Passes Unit Tests
  key: pytorch-compilation-passes-unit-tests
  timeout_in_minutes: 20
+  device: h200_18gb
  source_file_dependencies:
    - vllm/__init__.py
    - vllm/_aiter_ops.py
@@ -110,6 +112,7 @@ steps:
 - label: PyTorch Fullgraph Smoke Test
  key: pytorch-fullgraph-smoke-test
  timeout_in_minutes: 35
+  device: h200_18gb
  source_file_dependencies:
  - vllm/__init__.py
  - vllm/_aiter_ops.py
@@ -5,6 +5,7 @@ steps:
 - label: Quantization
  key: quantization
  timeout_in_minutes: 90
+  device: h200_18gb
  source_file_dependencies:
  - csrc/
  - vllm/model_executor/layers/quantization
@@ -42,6 +43,7 @@ steps:
 - label: Quantized Models Test
  key: quantized-models-test
  timeout_in_minutes: 60
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers/quantization
  - tests/models/quantization
@@ -5,6 +5,7 @@ steps:
 - label: Samplers Test
  key: samplers-test
  timeout_in_minutes: 75
+  device: h200_18gb
  source_file_dependencies:
  - vllm/model_executor/layers
  - vllm/sampling_metadata.py
@@ -98,7 +98,16 @@ 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 and causes PyTorch's
+# CUDACachingAllocator to crash with an NVML assertion. Fall back to
+# CUDA runtime checks on MIG.
+_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.
@@ -167,6 +167,12 @@ 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."""
+        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):