[Misc] Fix various vLLM import issues (#900)

Signed-off-by: Jacob Platin <[email protected]>
Co-authored-by: Jacob Platin <[email protected]>

Author: jrplatin

tpu_inference/models/common/model_loader.py

@@ -9,7 +9,7 @@
 from torchax.ops.mappings import j2t_dtype
 from transformers import PretrainedConfig
 from vllm.config import VllmConfig
-from vllm.utils import supports_kw
+from vllm.utils.functools import supports_kw
 
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils.quantization.quantization_utils import (

tpu_inference/platforms/tpu_jax.py

@@ -122,18 +122,18 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 "VLLM_ENABLE_V1_MULTIPROCESSING must be 0 when using Pathways(JAX_PLATFORMS=proxy)"
             )
 
-        from vllm.config import CompilationLevel
+        from vllm.config import CompilationMode
 
         cache_config = vllm_config.cache_config
         # For v0, the default block size is 16.
         if cache_config and cache_config.block_size is None:
             cache_config.block_size = cast(BlockSize, 16)
         compilation_config = vllm_config.compilation_config
 
-        # TPU only supports DYNAMO_ONCE compilation level
+        # TPU only supports DYNAMO_TRACE_ONCE compilation level
         # NOTE(xiang): the compilation_config is not used by jax.
-        if compilation_config.level != CompilationLevel.DYNAMO_ONCE:
-            compilation_config.level = CompilationLevel.DYNAMO_ONCE
+        if compilation_config.level != CompilationMode.DYNAMO_TRACE_ONCE:
+            compilation_config.level = CompilationMode.DYNAMO_TRACE_ONCE
 
         if compilation_config.backend == "":
             compilation_config.backend = "openxla"

tpu_inference/runner/input_batch_jax.py

@@ -9,7 +9,7 @@
 import numpy as np
 from vllm.lora.request import LoRARequest
 from vllm.sampling_params import SamplingType
-from vllm.utils import swap_dict_values
+from vllm.utils.collections import swap_dict_values
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.spec_decode.utils import is_spec_decode_unsupported