Apply input processor

vllm/engine/async_llm_engine.py

@@ -265,9 +265,11 @@ async def process_model_inputs_async(
         else:
             prompt_token_ids = inputs["prompt_token_ids"]
 
-        return LLMInputs(prompt_token_ids=prompt_token_ids,
-                         prompt=inputs.get("prompt"),
-                         multi_modal_data=inputs.get("multi_modal_data"))
+        llm_inputs = LLMInputs(prompt_token_ids=prompt_token_ids,
+                               prompt=inputs.get("prompt"),
+                               multi_modal_data=inputs.get("multi_modal_data"))
+
+        return self.input_processor(llm_inputs)
 
     async def add_request_async(
         self,

vllm/engine/llm_engine.py

@@ -21,7 +21,7 @@
 from vllm.engine.output_processor.util import create_output_by_sequence_group
 from vllm.executor.executor_base import ExecutorBase
 from vllm.executor.ray_utils import initialize_ray_cluster
-from vllm.inputs import LLMInputs, PromptInputs
+from vllm.inputs import INPUT_REGISTRY, LLMInputs, PromptInputs
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.outputs import (EmbeddingRequestOutput, RequestOutput,
@@ -219,6 +219,9 @@ def __init__(
         self.generation_config_fields = _load_generation_config_dict(
             model_config)
 
+        self.input_processor = INPUT_REGISTRY.create_input_processor(
+            self.model_config)
+
         self.model_executor = executor_class(
             model_config=model_config,
             cache_config=cache_config,
@@ -484,9 +487,11 @@ def process_model_inputs(
         else:
             prompt_token_ids = inputs["prompt_token_ids"]
 
-        return LLMInputs(prompt_token_ids=prompt_token_ids,
-                         prompt=inputs.get("prompt"),
-                         multi_modal_data=inputs.get("multi_modal_data"))
+        llm_inputs = LLMInputs(prompt_token_ids=prompt_token_ids,
+                               prompt=inputs.get("prompt"),
+                               multi_modal_data=inputs.get("multi_modal_data"))
+
+        return self.input_processor(llm_inputs)
 
     def add_request(
         self,

vllm/inputs/registry.py

@@ -1,9 +1,10 @@
+import functools
 from typing import (TYPE_CHECKING, Callable, Dict, Optional, Tuple, Type,
                     TypeVar)
-from typing_extensions import Concatenate, ParamSpec
 
 from torch import nn
 from transformers import PretrainedConfig
+from typing_extensions import Concatenate, ParamSpec
 
 from vllm.logger import init_logger
 
@@ -31,6 +32,7 @@
 
 
 def _for_hf(hf_config_type: Type[C]):
+
     def wrapper(
         fn: Callable[Concatenate[C, P], R],
     ) -> Callable[Concatenate["ModelConfig", P], R]:
@@ -54,6 +56,7 @@ def inner(
 
 
 def _for_multimodal_hf(hf_config_type: Type[C]):
+
     def wrapper(
         factory: Callable[Concatenate["VisionLanguageConfig", C, P], R],
     ) -> Callable[Concatenate["ModelConfig", P], R]:
@@ -138,8 +141,8 @@ def for_hf(cls, hf_config_type: Type[C]):
         """
 
         def wrapper(
-            processor: Callable[[C, LLMInputs], LLMInputs],
-        ) -> InputProcessor:
+            processor: Callable[[C, LLMInputs],
+                                LLMInputs], ) -> InputProcessor:
             return _for_hf(hf_config_type)(processor)
 
         return wrapper
@@ -219,11 +222,8 @@ def wrapper(model_cls: N) -> N:
 
         return wrapper
 
-    def dummy_data_for_profiling(
-        self,
-        model_config: "ModelConfig",
-        seq_len: int,
-    ):
+    def dummy_data_for_profiling(self, model_config: "ModelConfig",
+                                 seq_len: int):
         """Create dummy data for memory profiling."""
         # Avoid circular import
         from vllm.model_executor.model_loader import get_model_architecture
@@ -279,3 +279,10 @@ def process_input(self, model_config: "ModelConfig",
                            f"model class {model_cls.__name__}.")
 
         return processor(model_config, inputs)
+
+    def create_input_processor(self, model_config: ModelConfig):
+        """
+        Create an input processor (see :meth:`process_input`) for a
+        specific model.
+        """
+        return functools.partial(self.process_input, model_config=model_config)