jdkansjn

Author: fin-griffin

sweep/infer.py

@@ -25,10 +25,14 @@
 import sys
 from typing import Any
 
-# Must be set before vllm is imported — vllm reads these at import time.
-# v0 engine has stable LoRA + tensor-parallel support; v1 does not.
+# Must be set before torch/vllm are imported — read at import/init time.
+# Disable torch.compile globally: vLLM's vocab_parallel_embedding has a
+# @torch.compile'd function (get_masked_input_and_mask) that invokes Triton,
+# which tries to compile cuda_utils.c via gcc and needs python3.12-dev headers
+# that are not installed on this server.  TORCHDYNAMO_DISABLE=1 makes all
+# @torch.compile decorators fall back to eager execution.
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
 # spawn (not fork) avoids CUDA/NCCL re-init issues across worker processes.
-os.environ["VLLM_USE_V1"] = "0"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import torch
@@ -99,10 +103,6 @@ def vllm_generate(
         dtype="bfloat16",
         max_model_len=2048,
         tensor_parallel_size=n_gpus,
-        # Disable CUDA graph capture and torch.compile/inductor AOT compilation.
-        # The server lacks python3.12-dev headers, which Triton needs to compile
-        # cuda_utils.c; enforce_eager bypasses that compilation entirely.
-        enforce_eager=True,
     )
 
     lora_request = LoRARequest("adapter", 1, checkpoint)