Remove SKIP_JAX_PRECOMPILE

- Same functionality can be achieve with vllm argument --enforce-eager
- It is better to remove duplicate configs to avoid confusion from users

Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>

Author: kyuyeunk

.buildkite/models/google_gemma-3-27b-it.yml

@@ -9,7 +9,7 @@ steps:
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
-          bash -c 'SKIP_JAX_PRECOMPILE=1 VLLM_XLA_CHECK_RECOMPILATION=0 python3 /workspace/tpu_inference/examples/offline_inference.py --model=google/gemma-3-27b-it --tensor_parallel_size=8 --task=generate --max_model_len=1024 --max_num_seqs=1'
+          bash -c 'VLLM_XLA_CHECK_RECOMPILATION=0 python3 /workspace/tpu_inference/examples/offline_inference.py --model=google/gemma-3-27b-it --tensor_parallel_size=8 --task=generate --max_model_len=1024 --max_num_seqs=1'
   - label: "Record unit test result for google/gemma-3-27b-it"
     key: "record_google_gemma-3-27b-it_UnitTest"
     depends_on: "google_gemma-3-27b-it_UnitTest"

examples/disagg/run_disagg_multi_host.sh

@@ -63,7 +63,6 @@ for ((i=0; i<NUM_HOSTS_PER_INSTANCE; i++)); do
         -e TPU_KV_TRANSFER_PORT="${KV_PORT}" \
         -e TPU_SIDE_CHANNEL_PORT="${SIDE_PORT}" \
         -e RAY_DEDUP_LOGS="0" \
-        -e SKIP_JAX_PRECOMPILE="1" \
         \
         -e TPU_CHIPS_PER_PROCESS_BOUNDS="1,1,1" \
         -e TPU_PROCESS_BOUNDS="2,2,1" \
@@ -95,6 +94,7 @@ docker exec node-0 /bin/bash -c \
     --gpu-memory-utilization 0.3 \
     --tensor-parallel-size 4 \
     --kv-transfer-config '{\"kv_connector\":\"TPUConnector\",\"kv_connector_module_path\":\"tpu_inference.distributed.tpu_connector\",\"kv_role\":\"kv_producer\"}' \
+    --enforce-eager \
     > /root/logs/prefill.txt 2>&1 &"
 set +x
 
@@ -137,7 +137,6 @@ for ((i=0; i<NUM_HOSTS_PER_INSTANCE; i++)); do
         -e TPU_KV_TRANSFER_PORT="${KV_PORT}" \
         -e TPU_SIDE_CHANNEL_PORT="${SIDE_PORT}" \
         -e RAY_DEDUP_LOGS="0" \
-        -e SKIP_JAX_PRECOMPILE="1" \
         \
         -e TPU_CHIPS_PER_PROCESS_BOUNDS="1,1,1" \
         -e TPU_PROCESS_BOUNDS="2,2,1" \
@@ -169,5 +168,6 @@ docker exec node-20 /bin/bash -c \
     --gpu-memory-utilization 0.3 \
     --tensor-parallel-size 4 \
     --kv-transfer-config '{\"kv_connector\":\"TPUConnector\",\"kv_connector_module_path\":\"tpu_inference.distributed.tpu_connector\",\"kv_role\":\"kv_consumer\"}' \
+    --enforce-eager \
     > /root/logs/decode.txt 2>&1 &"
 set +x

examples/disagg/run_disagg_single_host.sh

@@ -45,13 +45,13 @@ for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
     \
     TPU_KV_TRANSFER_PORT=$KV_PORT \
     TPU_SIDE_CHANNEL_PORT=$SIDE_PORT \
-    SKIP_JAX_PRECOMPILE=1 \
     \
     vllm serve $MODEL \
     --port $PORT \
     --gpu-memory-utilization 0.2 \
     --tensor-parallel-size $PREFILLER_TP_SIZE \
     --kv-transfer-config "{\"kv_connector\":\"TPUConnector\",\"kv_connector_module_path\":\"tpu_inference.distributed.tpu_connector\",\"kv_role\":\"kv_producer\"}" \
+    --enforce-eager \
     > $HOME/logs/prefill_$i.txt 2>&1 &
 
     PREFILL_HOSTS+=("localhost")
@@ -72,13 +72,13 @@ for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
     \
     TPU_KV_TRANSFER_PORT=$KV_PORT \
     TPU_SIDE_CHANNEL_PORT=$SIDE_PORT \
-    SKIP_JAX_PRECOMPILE=1 \
     \
     vllm serve $MODEL \
     --port $PORT \
     --gpu-memory-utilization 0.2 \
     --tensor-parallel-size $DECODER_TP_SIZE \
     --kv-transfer-config "{\"kv_connector\":\"TPUConnector\",\"kv_connector_module_path\":\"tpu_inference.distributed.tpu_connector\",\"kv_role\":\"kv_consumer\"}" \
+    --enforce-eager \
     > $HOME/logs/decode_$i.txt 2>&1 &
 
     DECODE_HOSTS+=("localhost")

examples/offline_inference.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
-
 import vllm.envs as envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -17,6 +15,9 @@ def create_parser():
     parser.set_defaults(model="meta-llama/Llama-3.2-1B-Instruct")
     parser.set_defaults(max_model_len=1024)
 
+    # Skip long warmup for local simple test.
+    parser.set_defaults(enforce_eager=True)
+
     # Add sampling params
     sampling_group = parser.add_argument_group("Sampling parameters")
     sampling_group.add_argument("--max-tokens", type=int)
@@ -103,9 +104,6 @@ def main(args: dict):
 
 
 if __name__ == "__main__":
-    # Skip long warmup for local simple test.
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
-
     parser = create_parser()
     args: dict = vars(parser.parse_args())
 

examples/offline_lora_inference.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 import time
 
 import vllm.envs as envs
@@ -20,6 +19,9 @@ def create_parser():
     parser.set_defaults(enable_lora=True)
     parser.set_defaults(max_lora_rank=8)
 
+    # Skip long warmup for local simple test.
+    parser.set_defaults(enforce_eager=True)
+
     # Add sampling params
     sampling_group = parser.add_argument_group("Sampling parameters")
     sampling_group.add_argument("--max-tokens", type=int, default=16)
@@ -76,9 +78,6 @@ def main(args: dict):
 
 
 if __name__ == "__main__":
-    # Skip long warmup for local simple test.
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
-
     parser = create_parser()
     args: dict = vars(parser.parse_args())
 

tests/e2e/benchmarking/mm_bench.sh

@@ -91,7 +91,7 @@ checkThroughputAndRouge() {
 }
 
 echo "Spinning up the vLLM server..."
-(SKIP_JAX_PRECOMPILE=1 VLLM_XLA_CHECK_RECOMPILATION=0 vllm serve "$model_name" --max-model-len "$max_model_len" --max-num-seqs "$max_num_seqs" --disable-log-requests --max-num-batched-tokens "$max_batched_tokens" 2>&1 | tee -a "$LOG_FILE") &
+(VLLM_XLA_CHECK_RECOMPILATION=0 vllm serve "$model_name" --max-model-len "$max_model_len" --max-num-seqs "$max_num_seqs" --disable-log-requests --max-num-batched-tokens "$max_batched_tokens" --enforce-eager 2>&1 | tee -a "$LOG_FILE") &
 
 
 # Run a busy loop to block until the server is ready to receive requests

tests/e2e/test_data_parallel.py

@@ -67,6 +67,7 @@ def _run_inference_with_config(model_name: str,
         additional_config=additional_config,
         kv_cache_dtype=kv_cache_dtype,
         async_scheduling=async_scheduling,
+        enforce_eager=True,
     )
 
     engine_args_dict = asdict(engine_args)
@@ -173,7 +174,6 @@ def test_data_parallelism_correctness(
     This test compares outputs from a single-device run with data parallel runs
     to ensure correctness, including log probabilities.
     """
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'
     os.environ['VLLM_XLA_CHECK_RECOMPILATION'] = '0'
     model_name = "Qwen/Qwen2.5-1.5B-Instruct"
     # Use a smaller subset of prompts for correctness testing

tests/e2e/test_multi_modal_inference.py

@@ -24,7 +24,6 @@ def test_multi_modal_inference(monkeypatch):
     """
     Runs multi-modal inference and verifies the output.
     """
-    os.environ['SKIP_JAX_PRECOMPILE'] = '1'  # Skip warmup to save time.
     os.environ[
         'VLLM_XLA_CHECK_RECOMPILATION'] = '0'  # Allow compilation during execution.
 
@@ -65,6 +64,7 @@ def test_multi_modal_inference(monkeypatch):
             "fps": 1,
         },
         limit_mm_per_prompt={modality: 1},
+        enforce_eager=True,  # Skip warmup to save time.
     )
     engine_args = asdict(engine_args)
     llm = LLM(**engine_args)

tests/test_envs.py

@@ -56,13 +56,6 @@ def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
 
 
 def test_boolean_env_vars(monkeypatch: pytest.MonkeyPatch):
-    # Test SKIP_JAX_PRECOMPILE (default False)
-    assert envs.SKIP_JAX_PRECOMPILE is False
-    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "1")
-    assert envs.SKIP_JAX_PRECOMPILE is True
-    monkeypatch.setenv("SKIP_JAX_PRECOMPILE", "0")
-    assert envs.SKIP_JAX_PRECOMPILE is False
-
     # Test NEW_MODEL_DESIGN (default False)
     assert envs.NEW_MODEL_DESIGN is False
     monkeypatch.setenv("NEW_MODEL_DESIGN", "1")
@@ -133,7 +126,6 @@ def test_dir_returns_all_env_vars():
     assert len(env_vars) == len(environment_variables)
     assert "JAX_PLATFORMS" in env_vars
     assert "TPU_NAME" in env_vars
-    assert "SKIP_JAX_PRECOMPILE" in env_vars
     assert "MODEL_IMPL_TYPE" in env_vars
 
 

tpu_inference/envs.py

@@ -14,7 +14,6 @@
     TPU_MULTIHOST_BACKEND: str = ""
     PREFILL_SLICES: str = ""
     DECODE_SLICES: str = ""
-    SKIP_JAX_PRECOMPILE: bool = False
     MODEL_IMPL_TYPE: str = "flax_nnx"
     NEW_MODEL_DESIGN: bool = False
     PHASED_PROFILING_DIR: str = ""
@@ -45,9 +44,6 @@
     # Slice configuration for disaggregated decode workers
     "DECODE_SLICES":
     lambda: os.getenv("DECODE_SLICES", ""),
-    # Skip JAX precompilation step during initialization
-    "SKIP_JAX_PRECOMPILE":
-    lambda: bool(int(os.getenv("SKIP_JAX_PRECOMPILE", "0"))),
     # Model implementation type (e.g., "flax_nnx")
     "MODEL_IMPL_TYPE":
     lambda: os.getenv("MODEL_IMPL_TYPE", "flax_nnx").lower(),