[CI] Fixes to catchup with vllm changes (#912)

Signed-off-by: Hongmin Fan <[email protected]>

Author: hfan

tests/layers/vllm/test_unquantized.py

@@ -1,3 +1,4 @@
+import os
 import tempfile
 
 import jax
@@ -415,6 +416,7 @@ def test_merged_column_parallel_linear(model, bias, mesh, fuse_matmuls,
 @pytest.mark.parametrize("topk", [2])
 def test_fused_moe(use_ep, mesh, num_tokens, intermediate_size, hidden_size,
                    num_experts, topk):
+    os.environ['VLLM_DISABLE_SHARED_EXPERTS_STREAM'] = '1'
     torch.manual_seed(42)
     dtype = torch.bfloat16
 

tpu_inference/core/disagg_executor.py

@@ -6,8 +6,9 @@
 from vllm.logger import init_logger
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.cache import worker_receiver_cache_from_config
-from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
-                        run_method)
+from vllm.utils import run_method
+from vllm.utils.network_utils import (get_distributed_init_method, get_ip,
+                                      get_open_port)
 from vllm.v1.executor.abstract import Executor
 from vllm.v1.outputs import AsyncModelRunnerOutput
 from vllm.v1.worker.worker_base import WorkerWrapperBase

tpu_inference/distributed/utils.py

@@ -1,6 +1,6 @@
 import os
 
-from vllm.utils import get_ip
+from vllm.utils.network_utils import get_ip
 
 from tpu_inference.logger import init_logger
 

tpu_inference/executors/ray_distributed_executor.py

@@ -1,18 +1,22 @@
 import os
-from typing import Dict, List, Optional
+from array import array
+from typing import Any, Dict, List, Optional
 
 import ray
 import vllm.envs as envs
 from ray.util.placement_group import PlacementGroup
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
-from vllm.executor.ray_distributed_executor import RayWorkerMetaData
-from vllm.executor.ray_utils import RayWorkerWrapper, _wait_until_pg_ready
+from vllm.multimodal.inputs import MultiModalKwargs
 from vllm.platforms import current_platform
 from vllm.ray.ray_env import get_env_vars_to_copy
-from vllm.utils import get_distributed_init_method, get_ip, get_open_port
+from vllm.sequence import VLLM_TOKEN_ID_ARRAY_TYPE
+from vllm.utils.network_utils import (get_distributed_init_method, get_ip,
+                                      get_open_port)
 from vllm.v1.executor.ray_distributed_executor import \
     RayDistributedExecutor as RayDistributedExecutorV1
+from vllm.v1.executor.ray_executor import RayWorkerMetaData
+from vllm.v1.executor.ray_utils import RayWorkerWrapper, _wait_until_pg_ready
 
 from tpu_inference.logger import init_logger
 
@@ -27,14 +31,27 @@
 from collections import defaultdict
 
 import msgspec
-from vllm.executor.msgspec_utils import encode_hook
 from vllm.v1.outputs import SamplerOutput
 
 from tpu_inference.distributed.utils import set_node_kv_ip_port
 
 logger = init_logger(__name__)
 
 
+def _encode_hook(obj: Any) -> Any:
+    """Custom msgspec enc hook that supports array types and MultiModalKwargs.
+
+    See https://jcristharif.com/msgspec/api.html#msgspec.msgpack.Encoder
+    """
+    if isinstance(obj, array):
+        assert obj.typecode == VLLM_TOKEN_ID_ARRAY_TYPE, (
+            f"vLLM array type should use '{VLLM_TOKEN_ID_ARRAY_TYPE}' type. "
+            f"Given array has a type code of {obj.typecode}.")
+        return obj.tobytes()
+    if isinstance(obj, MultiModalKwargs):
+        return dict(obj)
+
+
 class RayDistributedExecutor(RayDistributedExecutorV1):
     """Ray-based distributed executor for TPU.
 
@@ -82,7 +99,7 @@ def _init_executor(self) -> None:
         # Create the parallel GPU workers.
         self._init_workers_ray(placement_group)
 
-        self.input_encoder = msgspec.msgpack.Encoder(enc_hook=encode_hook)
+        self.input_encoder = msgspec.msgpack.Encoder(enc_hook=_encode_hook)
         self.output_decoder = msgspec.msgpack.Decoder(
             Optional[List[SamplerOutput]])
         self.use_v1 = envs.VLLM_USE_V1

tpu_inference/layers/vllm/quantization/awq.py

@@ -11,12 +11,11 @@
 from vllm.model_executor.layers.quantization import \
     register_quantization_config
 from vllm.model_executor.layers.quantization.awq import (AWQConfig,
-                                                         AWQLinearMethod,
-                                                         is_layer_skipped_awq)
+                                                         AWQLinearMethod)
 from vllm.model_executor.layers.quantization.base_config import \
     QuantizeMethodBase
-from vllm.model_executor.layers.quantization.utils.quant_utils import \
-    unpack_quantized_values_into_int32
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    is_layer_skipped, unpack_quantized_values_into_int32)
 from vllm.scalar_type import scalar_types
 
 from tpu_inference.layers.vllm.linear_common import (
@@ -48,7 +47,7 @@ def get_quant_method(
     ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:
         if isinstance(layer, LinearBase):
             linear_config = self.get_linear_config(layer)
-            if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
+            if is_layer_skipped(prefix, self.modules_to_not_convert):
                 return VllmUnquantizedLinearMethod(linear_config)
             return VllmAWQLinearMethod(self, linear_config)
         elif isinstance(layer, FusedMoE):

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.func_utils import supports_kw
 
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils.quantization.quantization_utils import (

tpu_inference/models/vllm/vllm_model_wrapper.py

@@ -40,24 +40,15 @@ def __init__(self, vllm_model: torch.nn.Module):
         self.vllm_model = vllm_model
 
     def forward(self, **kwargs) -> torch.Tensor:
-        # We don't support multimodal input in Gemma3, but we need patch it to
-        # None to workaround vLLM Gemma3 model bug that
-        # `get_multimodal_embeddings` returns empty list but it's caller checks
-        # for None.
-        with patch(
-                "vllm.model_executor.models.gemma3_mm."
-                "Gemma3ForConditionalGeneration."
-                "get_multimodal_embeddings",
-                return_value=None):
-            if "hidden_state" in kwargs:
-                return self.compute_logits(kwargs["hidden_state"])
-            else:
-                return self.compute_hidden_state(
-                    kwargs["input_ids"],
-                    kwargs["positions"],
-                    kwargs["intermediate_tensors"],
-                    kwargs["inputs_embeds"],
-                )
+        if "hidden_state" in kwargs:
+            return self.compute_logits(kwargs["hidden_state"])
+        else:
+            return self.compute_hidden_state(
+                kwargs["input_ids"],
+                kwargs["positions"],
+                kwargs["intermediate_tensors"],
+                kwargs["inputs_embeds"],
+            )
 
     def compute_hidden_state(
         self,

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.mode != CompilationMode.DYNAMO_TRACE_ONCE:
+            compilation_config.mode = 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.collection_utils import swap_dict_values
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.spec_decode.utils import is_spec_decode_unsupported