Rebase to vLLM 0.14.1

arctic_inference/vllm/args.py

@@ -20,7 +20,7 @@
 
 from vllm.config import ParallelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
-from vllm.utils import FlexibleArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from arctic_inference.patching import ArcticPatch
 from arctic_inference.vllm.config import ArcticParallelConfig
@@ -50,7 +50,6 @@ class EngineArgsPatch(ArcticPatch[EngineArgs]):
     _orig_add_cli_args = EngineArgs.add_cli_args
     _orig_from_cli_args = EngineArgs.__dict__["from_cli_args"].__wrapped__
     _orig_create_engine_config = EngineArgs.create_engine_config
-    _orig_is_v1_supported_oracle = EngineArgs._is_v1_supported_oracle
 
     def __new__(cls, *args, **kwargs):
         # Override __new__ to return an ArcticEngineArgs instead of an
@@ -109,6 +108,11 @@ def create_engine_config(self, *args, **kwargs):
         if (self.ulysses_sequence_parallel_size > 1 and
                 self.distributed_executor_backend is None):
             self.distributed_executor_backend = "mp"
+
+        # Store ulysses_sequence_parallel_size for access during config initialization
+        from arctic_inference.vllm import ulysses
+        ulysses._ulysses_sp_size = self.ulysses_sequence_parallel_size
+
         vllm_config = self._orig_create_engine_config(*args, **kwargs)
         # Recreate the parallel config with Arctic parameters since they might
         # not be passed to the parallel config __init__ when first initialized.
@@ -121,21 +125,6 @@ def create_engine_config(self, *args, **kwargs):
         vllm_config.parallel_config = ArcticParallelConfig(**kwargs)
         return vllm_config
 
-    def _is_v1_supported_oracle(self, *args, **kwargs):
-        orig_speculative_config = self.speculative_config
-
-        # Since Arctic Inference is only compatible with v1 and we already
-        # check it earlier, we can just disable this check altogether.
-        if (self.speculative_config is not None and
-                self.speculative_config.get("method") in ("arctic", "suffix")):
-            self.speculative_config = None
-
-        res = self._orig_is_v1_supported_oracle(*args, **kwargs)
-
-        self.speculative_config = orig_speculative_config
-
-        return res
-
 
 class AsyncEngineArgsPatch(ArcticPatch[AsyncEngineArgs]):
 

arctic_inference/vllm/config.py

@@ -13,9 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
+from pydantic.dataclasses import dataclass
 import logging
 
+import vllm
 from vllm.config import ParallelConfig, SpeculativeConfig, VllmConfig
 from vllm.transformers_utils.configs.mlp_speculator import MLPSpeculatorConfig
 
@@ -55,6 +56,7 @@ def world_size(self, value: int) -> None:
 @dataclass
 class ArcticSpeculativeConfig(SpeculativeConfig):
 
+    method: str | None = None
     enable_suffix_decoding: bool = False
     suffix_cache_max_depth: int = 64
     suffix_speculative_tokens: int = 0
@@ -80,21 +82,18 @@ class SpeculativeConfigPatch(ArcticPatch[SpeculativeConfig]):
     _orig_post_init = SpeculativeConfig.__post_init__
 
     def __new__(cls, *args, **kwargs):
-        # Override __new__ to return an ArcticSpeculativeConfig instead of a
-        # SpeculativeConfig when creating a new instance of the class.
         if cls is SpeculativeConfig:
             return ArcticSpeculativeConfig.__new__(ArcticSpeculativeConfig,
                                                    *args, **kwargs)
         return super(SpeculativeConfig, cls).__new__(cls)
 
     def __post_init__(self):
-        use_suffix = (self.method
-                      == "suffix") or (self.method is None
-                                       and self.enable_suffix_decoding)
-        use_hybrid = (self.method == "arctic"
-                      and self.enable_suffix_decoding)
-        if (use_suffix or self.method == "arctic") and \
-            self.disable_by_batch_size is None:
+        is_arctic_method = self.method in ("arctic", "mlp_speculator")
+        use_suffix = (self.method == "suffix") or (self.method is None 
+                                                   and self.enable_suffix_decoding)
+        use_hybrid = (self.method == "arctic" and self.enable_suffix_decoding)
+
+        if (use_suffix or is_arctic_method) and self.disable_by_batch_size is None:
             logger.info("Defaulting disable_by_batch_size to 64")
             self.disable_by_batch_size = 64
 
@@ -104,15 +103,42 @@ def __post_init__(self):
         if use_suffix:
             self.method = "suffix"
             self.enable_suffix_decoding = True
-            self.num_speculative_tokens = self.suffix_cache_max_depth
+            # Use suffix_speculative_tokens if explicitly set, otherwise
+            # default to 16 (not suffix_cache_max_depth which can be very
+            # large and makes every step process 1+N tokens even when the
+            # suffix cache has no matches).
+            # NOTE: num_speculative_tokens defaults to None (not 0).
+            if self.suffix_speculative_tokens > 0:
+                self.num_speculative_tokens = self.suffix_speculative_tokens
+            elif self.num_speculative_tokens is None:
+                self.num_speculative_tokens = 16
             self._verify_args()
+            return 
+
+        if is_arctic_method:
+            actual_draft_model = getattr(self, "draft_model", None)
+
+            self.draft_model = None 
+
+            try:
+                self._orig_post_init()
+            finally:
+                self.draft_model = actual_draft_model
+
+            if self.num_speculative_tokens == 0:
+                self.num_speculative_tokens = getattr(self, "num_lookahead_slots", 1)
         else:
             self._orig_post_init()
 
 
 class VllmConfigPatch(ArcticPatch[VllmConfig]):
 
     _orig_str = VllmConfig.__str__
+    _orig_post_init = VllmConfig.__post_init__
+
+    from typing import Literal
+    OldEagleModelTypes = vllm.config.speculative.EagleModelTypes
+    NewEagleModelTypes = Literal["arctic", "suffix", OldEagleModelTypes]
 
     def __str__(self, *args, **kwargs):
         string = self._orig_str(*args, **kwargs)
@@ -121,6 +147,24 @@ def __str__(self, *args, **kwargs):
         string += f", shift_parallel_threshold={self.parallel_config.shift_parallel_threshold}"
         return string
 
+    def __post_init__(self, *args, **kwargs):
+        # if self.speculative_config is not None:
+        #     if self.speculative_config.method not in get_args(EagleModelTypes):
+        #         raise ValueError(
+        #             "Currently, async scheduling is only supported "
+        #             "with EAGLE/MTP kind of speculative decoding"
+        #         )
+        import sys
+        from typing import Literal
+        target_module = sys.modules[VllmConfig.__module__]
+        original_types = getattr(target_module, "EagleModelTypes")
+        NewEagleModelTypes = Literal["mlp_speculator", "suffix", original_types]
+        setattr(target_module, "EagleModelTypes", NewEagleModelTypes)
+        try:
+            self._orig_post_init(*args, **kwargs)
+        finally:
+            setattr(target_module, "EagleModelTypes", original_types)
+
 
 class MLPSpeculatorConfigPatch(ArcticPatch[MLPSpeculatorConfig]):
 
@@ -129,3 +173,16 @@ class MLPSpeculatorConfigPatch(ArcticPatch[MLPSpeculatorConfig]):
     def __init__(self, *args, **kwargs):
         self.base_model_arch = kwargs.pop("base_model_arch", "")
         self._orig_init(*args, **kwargs)
+
+        # Inject dummy attributes required by vLLM's ModelArchConfigConvertor
+        # The convertor tries to calculate head_size = hidden_size // num_attention_heads
+        if not hasattr(self, "num_attention_heads"):
+            self.num_attention_heads = 1
+
+        if not hasattr(self, "hidden_size"):
+            # Fallback to n_embd if present, otherwise default to a safe dummy value
+            self.hidden_size = getattr(self, "n_embd", 1024)
+
+        # Ensure hidden_size is an integer to prevent TypeError during division
+        if hasattr(self, "hidden_size"):
+            self.hidden_size = int(self.hidden_size)