ignore type

Author: LiuXiaoxuanPKU

vllm/engine/arg_utils.py

@@ -1476,8 +1476,8 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
                 ("model" in self.speculative_config and
                  self.speculative_config["model"] in ("ngram", "[ngram]"))):
                 is_ngram_enabled = True
-            elif (("model" in self.speculative_config and 
-                   "eagle" in self.speculative_config["model"].lower())):
+            elif ("model" in self.speculative_config
+                  and "eagle" in self.speculative_config["model"].lower()):
                 is_eagle_enabled = True
             else:
                 _raise_or_fallback(feature_name="Speculative Decoding",
@@ -1517,7 +1517,7 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
 
         # LoRA is supported on V1, but off by default for now.
         if self.enable_lora and _warn_or_fallback("LORA"):
-            return Falsef
+            return False
 
         # PP is supported on V1 with Ray distributed executor,
         # but off for MP distributed executor for now.
@@ -1529,7 +1529,7 @@ def _is_v1_supported_oracle(self, model_config: ModelConfig) -> bool:
         # ngram is supported on V1, but off by default for now.
         if is_ngram_enabled and _warn_or_fallback("ngram"):
             return False
-        
+
         if is_eagle_enabled and _warn_or_fallback("eagle"):
             return False
 

vllm/v1/spec_decode/ngram_proposer.py

@@ -4,9 +4,14 @@
 import numpy as np
 from numba import jit
 
+from vllm.config import VllmConfig
+
 
 class NgramProposer:
 
+    def __init__(self, vllm_config: VllmConfig):
+        self.vllm_config = vllm_config
+
     def propose(
         self,
         context_token_ids: np.ndarray,

vllm/v1/worker/gpu_model_runner.py

@@ -159,9 +159,11 @@ def __init__(
             self.use_spec_decode = True
             if get_pp_group().is_last_rank:
                 if self.speculative_config.method == "ngram":
-                    self.drafter = NgramProposer(self.vllm_config)
+                    self.drafter = NgramProposer(
+                        self.vllm_config)  # type:ignore
                 elif self.speculative_config.method == "eagle":
-                    self.drafter = EagleProposer(self.vllm_config, self.device)
+                    self.drafter = EagleProposer(self.vllm_config,
+                                                 self.device)  # type:ignore
                 else:
                     raise ValueError("Unknown speculative decoding method: "
                                      f"{self.speculative_config.method}")
@@ -1143,9 +1145,11 @@ def execute_model(
             # Speculative decoding is not enabled.
             spec_token_ids = None
         elif self.speculative_config.method == "ngram":
+            assert isinstance(self.drafter, NgramProposer)
             spec_token_ids = self.generate_draft_token_ids(
                 valid_sampled_token_ids, sampling_metadata)
         elif self.speculative_config.method == "eagle":
+            assert isinstance(self.drafter, EagleProposer)
             # TODO(woosuk): Refactor the loop.
             next_token_ids: list[int] = []
             for i, token_ids in enumerate(valid_sampled_token_ids):
@@ -1265,7 +1269,7 @@ def load_model(self) -> None:
                                                   self.lora_config,
                                                   self.device)
             if (hasattr(self, "drafter")
-                    and self.speculative_config.method != "ngram"):
+                    and not isinstance(self.drafter, NgramProposer)):
                 logger.info("Loading drafter model...")
                 self.drafter.load_model(self.model)
             time_after_load = time.perf_counter()