Bump vllm to v0.10.1 and add compatibility code (#443)

Signed-off-by: Max de Bayser <[email protected]>

Author: maxdebayser

pyproject.toml

@@ -13,7 +13,7 @@ license = {text = "Apache 2"}
 dependencies = [
     "fms-model-optimizer[fp8]>=0.6.0",
     "ibm-fms>=1.2.1",
-    "vllm>=0.9.2,<=0.10.0",
+    "vllm>=0.9.2,<=0.10.1.1",
 ]
 requires-python = ">=3.9"
 dynamic = ["version"]

tests/utils/test_upstream_compatibility.py

@@ -168,3 +168,20 @@ def test_mm_inputs():
                                "renamed mm_inputs to mm_kwargs.")
         # The compat code introduced in the PR below can now be removed:
         # https://github.com/vllm-project/vllm-spyre/pull/380
+
+
+@pytest.mark.cpu
+def test_init_builtin_logitsprocs():
+
+    import vllm.v1.sample.logits_processor
+    has_init_builtin_logitsprocs = hasattr(vllm.v1.sample.logits_processor,
+                                           "init_builtin_logitsprocs")
+
+    if VLLM_VERSION == "vLLM:main":
+        assert not has_init_builtin_logitsprocs
+    elif VLLM_VERSION == "vLLM:lowest":
+        assert has_init_builtin_logitsprocs, (
+            "The lowest supported vLLM version already"
+            "refactored init_builtin_logitsprocs.")
+        # The compat code introduced in the PR below can now be removed:
+        # https://github.com/vllm-project/vllm-spyre/pull/443

vllm_spyre/v1/worker/spyre_input_batch.py

@@ -5,16 +5,17 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Generic, Optional, TypeVar, cast
+from typing import Any, Generic, Optional, TypeVar, cast
 
 import numpy as np
 import torch
+import vllm.v1.sample.logits_processor
+from vllm.config import VllmConfig
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams, SamplingType
 from vllm.v1.pool.metadata import PoolingMetadata
 from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,
-                                             MoveDirectionality,
-                                             init_builtin_logitsprocs)
+                                             MoveDirectionality)
 from vllm.v1.sample.metadata import SamplingMetadata
 
 
@@ -200,6 +201,29 @@ def num_tokens(self) -> int:
         return len(self.prompt_token_ids) + len(self.output_token_ids)
 
 
+# Compatibility code, remove when no supported version
+# has init_builtin_logitsprocs any more
+def get_builtin_logits_processors(
+        vllm_config: Optional[VllmConfig] = None) -> Any:
+    if hasattr(vllm.v1.sample.logits_processor, "LogitsProcessors"):
+        if vllm_config is None:
+            return vllm.v1.sample.logits_processor.LogitsProcessors()
+        return vllm.v1.sample.logits_processor.LogitsProcessors(
+            ctor(vllm_config, "cpu", False)
+            for ctor in vllm.v1.sample.logits_processor.\
+                BUILTIN_LOGITS_PROCESSORS)
+    else:
+        if vllm_config is None:
+            return vllm.v1.sample.logits_processor.LogitsProcessorManager(
+                non_argmax_invariant=[],
+                argmax_invariant=[],
+            )
+        return vllm.v1.sample.logits_processor.init_builtin_logitsprocs(
+            pin_memory_available=False,
+            max_num_reqs=vllm_config.scheduler_config.max_num_seqs + 1,
+            device="cpu")
+
+
 class SamplingInputBatch(BaseInputBatch[SamplingRequestState]):
     '''
     This class was based on the InputBatch for GPU of vLLM V1.
@@ -229,6 +253,8 @@ def __init__(
         device: torch.device,
         pin_memory: bool,
         vocab_size: int,
+        # Type here is any for compatibility reasons
+        logitsprocs: Optional[Any] = None,
     ):
         super().__init__(
             max_num_reqs,
@@ -297,13 +323,7 @@ def __init__(
         # updates. Should reset each step.
         self.batch_update_builder = BatchUpdateBuilder()
 
-        # Define logits processors.
-        # TODO(andy): logits processor list should be extensible via engine
-        # constructor argument; for now the list is fixed.
-        self.logitsprocs = init_builtin_logitsprocs(pin_memory_available=False,
-                                                    max_num_reqs=max_num_reqs +
-                                                    1,
-                                                    device=device)
+        self.logitsprocs = logitsprocs or get_builtin_logits_processors()
 
         self.has_allowed_token_ids: set[str] = set()
         self.allowed_token_ids_mask: Optional[torch.Tensor] = None

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -29,12 +29,9 @@
 from vllm_spyre.platform import SpyrePlatform
 # yapf conflicts with ruff for this block
 # yapf: disable
-from vllm_spyre.v1.worker.spyre_input_batch import (BaseInputBatch,
-                                                    BaseRequestState,
-                                                    PoolingInputBatch,
-                                                    PoolingRequestState,
-                                                    SamplingInputBatch,
-                                                    SamplingRequestState)
+from vllm_spyre.v1.worker.spyre_input_batch import (
+    BaseInputBatch, BaseRequestState, PoolingInputBatch, PoolingRequestState,
+    SamplingInputBatch, SamplingRequestState, get_builtin_logits_processors)
 
 # yapf: enable
 if TYPE_CHECKING:
@@ -306,12 +303,17 @@ def load_model(self, prompt_lens: Iterable[int],
         )
 
     def build_input_batch(self) -> SamplingInputBatch:
+        # Define logits processors.
+        # TODO(Max): logits processor list should be extensible via engine
+        # constructor argument; for now the list is fixed to builtin processors
+        logits_processors = get_builtin_logits_processors(self.vllm_config)
         return SamplingInputBatch(
             max_num_reqs=self.scheduler_config.max_num_seqs,
             max_model_len=self.model_config.max_model_len,
             device=self.device,
             pin_memory=self.pin_memory,
             vocab_size=self.model_config.get_vocab_size(),
+            logitsprocs=logits_processors,
         )
 
     @property
@@ -810,8 +812,7 @@ def __init__(
             max_model_len=vllm_config.model_config.max_model_len,
             device=self.device,
             pin_memory=self.pin_memory,
-            vocab_size=vllm_config.model_config.get_vocab_size(),
-        )
+            vocab_size=vllm_config.model_config.get_vocab_size())
 
     def pre_warmup(self) -> None:
         # Set the number of kv cache blocks to the minimal value of 2 which is
@@ -1351,9 +1352,18 @@ def build_input_batch(self) -> PoolingInputBatch:
     def load_model(self, prompt_lens: Iterable[int],
                    num_decode_tokens: Iterable[int]) -> None:
 
-        if self.model_config.task == "embed":
+        task = self.model_config.task
+        if task is None:
+            # Task is being deprecated upstream because the models
+            # support several tasks at once. But for now, here we need
+            # to know the task to load the model with
+            # AutoModelForSequenceClassification
+            task = self.model_config._get_default_pooling_task(
+                self.model_config.architectures)
+
+        if task == "embed":
             self.model = AutoModel.from_pretrained(self.model_config.model)
-        elif self.model_config.task == "classify":
+        elif task == "classify":
             class_model = AutoModelForSequenceClassification.from_pretrained(
                 self.model_config.model)
             if hasattr(class_model, "bert"):
@@ -1368,7 +1378,7 @@ def load_model(self, prompt_lens: Iterable[int],
                     "Bert or Roberta for sequence classification")
             self.classifier = class_model.classifier
         else:
-            raise ValueError(f"Unsupported task {self.model_config.task}")
+            raise ValueError(f"Unsupported task {task}")
 
         model_class_name = type(self.model).__name__
         self.is_roberta = "roberta" in model_class_name.lower()
@@ -1393,7 +1403,7 @@ def load_model(self, prompt_lens: Iterable[int],
                     dynamic=False,
                     backend=envs_spyre.VLLM_SPYRE_DYNAMO_BACKEND)
 
-        if self.model_config.task == "classify":
+        if task == "classify":
             tokenizer = AutoTokenizer.from_pretrained(self.model_config.model)
             output = tokenizer(text="foo", text_pair="bar")
             self.use_token_type_ids = "token_type_ids" in output
@@ -1404,13 +1414,13 @@ def load_model(self, prompt_lens: Iterable[int],
         if hasattr(Pooler, "from_config_with_defaults"):
             # TODO: remove this when we no longer support
             # vllm version v0.9.2
-            if self.model_config.task == "embed":
+            if task == "embed":
                 self.pooler = Pooler.from_config_with_defaults(
                     pooler_config,
                     pooling_type=PoolingType.CLS,
                     normalize=True,
                     softmax=False)
-            elif self.model_config.task == "classify":
+            elif task == "classify":
                 self.pooler = ClassifierPooler(config=self.model_config,
                                                pooler=self._pooler,
                                                classifier=self.classifier)
@@ -1428,10 +1438,10 @@ def load_model(self, prompt_lens: Iterable[int],
             if 'default_pooling_type' in annotations:
                 extra_args['default_pooling_type'] = PoolingType.CLS
 
-            if self.model_config.task == "embed":
+            if task == "embed":
                 self.pooler = Pooler.for_embed(pooler_config=pooler_config,
                                                **extra_args)
-            elif self.model_config.task == "classify":
+            elif task == "classify":
                 self.pooler = ClassifierPooler(
                     pooling=self._pooler,
                     classifier=self.classifier,