3fa2c12185 · tduncan/Nvidia-Optimization-Stack-VLLM

docs/design/io_processor_plugins.md

+14 -4

@@ -13,7 +13,6 @@ IOProcessorInput = TypeVar("IOProcessorInput")

13 IOProcessorOutput = TypeVar("IOProcessorOutput")

14

15 class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

16 -

16 def __init__(self, vllm_config: VllmConfig):

17 self.vllm_config = vllm_config

18

@@ -49,13 +48,24 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

48 request_id: str | None = None,

49 **kwargs,

50 ) -> IOProcessorOutput:

52 - collected_output = [item async for i, item in model_output]

51 + # We cannot guarantee outputs are returned in the same order they were

52 + # fed to vLLM.

53 + # Let's sort them by id before post_processing

54 + sorted_output = sorted(

55 + [(i, item) async for i, item in model_output], key=lambda output: output[0]

56 + )

57 + collected_output = [output[1] for output in sorted_output]

58 return self.post_process(collected_output, request_id, **kwargs)

59

60 @abstractmethod

61 def parse_request(self, request: Any) -> IOProcessorInput:

62 raise NotImplementedError

63

64 + def validate_or_generate_params(

65 + self, params: SamplingParams | PoolingParams | None = None

66 + ) -> SamplingParams | PoolingParams:

67 + return params or PoolingParams()

68 +

69 @abstractmethod

70 def output_to_response(

71 self, plugin_output: IOProcessorOutput

@@ -66,10 +76,10 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

76 The `parse_request` method is used for validating the user prompt and converting it into the input expected by the `pre_process`/`pre_process_async` methods.

77 The `pre_process*` methods take the validated plugin input to generate vLLM's model prompts for regular inference.

78 The `post_process*` methods take `PoolingRequestOutput` objects as input and generate a custom plugin output.

69 -

79 +The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.

80 The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/openai/serving_pooling.py).

81

72 -An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/christian-pinto/prithvi_io_processor_plugin). Please, also refer to our online ([examples/online_serving/prithvi_geospatial_mae.py](../../examples/online_serving/prithvi_geospatial_mae.py)) and offline ([examples/offline_inference/prithvi_geospatial_mae_io_processor.py](../../examples/offline_inference/prithvi_geospatial_mae_io_processor.py)) inference examples.

82 +An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/online_serving/prithvi_geospatial_mae.py](../../examples/online_serving/prithvi_geospatial_mae.py)) and offline ([examples/offline_inference/prithvi_geospatial_mae_io_processor.py](../../examples/offline_inference/prithvi_geospatial_mae_io_processor.py)) inference examples.

83

84 ## Using an IO Processor plugin

85

examples/offline_inference/prithvi_geospatial_mae.py

+1 -1

@@ -64,7 +64,7 @@ class PrithviMAE:

64 }

65

66 prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}

67 - outputs = self.model.encode(prompt, use_tqdm=False)

67 + outputs = self.model.encode(prompt, pooling_task="plugin", use_tqdm=False)

68

69 return outputs[0].outputs.data

70

examples/offline_inference/prithvi_geospatial_mae_io_processor.py

+5 -9

@@ -6,14 +6,14 @@ import os

6 import torch

7

8 from vllm import LLM

9 -from vllm.pooling_params import PoolingParams

9

10 # This example shows how to perform an offline inference that generates

11 # multimodal data. In this specific case this example will take a geotiff

12 # image as input, process it using the multimodal data processor, and

13 # perform inference.

15 -# Requirement - install plugin at:

16 -# https://github.com/christian-pinto/prithvi_io_processor_plugin

14 +# Requirements:

15 +# - install TerraTorch v1.1 (or later):

16 +# pip install terratorch>=v1.1

17

18

19 def main():

@@ -36,16 +36,12 @@ def main():

36 # to avoid the model going OOM.

37 # The maximum number depends on the available GPU memory

38 max_num_seqs=32,

39 - io_processor_plugin="prithvi_to_tiff",

39 + io_processor_plugin="terratorch_segmentation",

40 model_impl="terratorch",

41 enable_mm_embeds=True,

42 )

43

44 - pooling_params = PoolingParams(task="token_classify", activation=False)

45 - pooler_output = llm.encode(

46 - img_prompt,

47 - pooling_params=pooling_params,

48 - )

44 + pooler_output = llm.encode(img_prompt, pooling_task="plugin")

45 output = pooler_output[0].outputs

46

47 print(output)

examples/online_serving/prithvi_geospatial_mae.py

+3 -4

@@ -11,14 +11,14 @@ import requests

11 # image as input, process it using the multimodal data processor, and

12 # perform inference.

13 # Requirements :

14 -# - install plugin at:

15 -# https://github.com/christian-pinto/prithvi_io_processor_plugin

14 +# - install TerraTorch v1.1 (or later):

15 +# pip install terratorch>=v1.1

16 # - start vllm in serving mode with the below args

17 # --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'

18 # --model-impl terratorch

19 # --task embed --trust-remote-code

20 # --skip-tokenizer-init --enforce-eager

21 -# --io-processor-plugin prithvi_to_tiff

21 +# --io-processor-plugin terratorch_segmentation

22 # --enable-mm-embeds

23

24

@@ -35,7 +35,6 @@ def main():

35 },

36 "priority": 0,

37 "model": "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",

38 - "softmax": False,

38 }

39

40 ret = requests.post(server_endpoint, json=request_payload_url)

tests/models/multimodal/pooling/test_prithvi_mae.py

+1 -1

@@ -40,7 +40,7 @@ def _run_test(

40 max_num_seqs=32,

41 default_torch_num_threads=1,

42 ) as vllm_model:

43 - vllm_model.llm.encode(prompt, pooling_task="token_classify")

43 + vllm_model.llm.encode(prompt, pooling_task="plugin")

44

45

46 MODELS = ["mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py

+2 -2

@@ -368,9 +368,9 @@ class PrithviMultimodalDataProcessor(IOProcessor):

368 out_format = "b64_json"

369

370 for output in model_output:

371 - y_hat = output.outputs.data.argmax(dim=1)

371 + y_hat = output.outputs.data.argmax(dim=0)

372 pred = torch.nn.functional.interpolate(

373 - y_hat.unsqueeze(1).float(),

373 + y_hat[None, None, ...].float(),

374 size=self.img_size,

375 mode="nearest",

376 )

tests/plugins_tests/test_io_processor_plugins.py

+1 -6

@@ -9,7 +9,6 @@ from tests.utils import RemoteOpenAIServer

9 from vllm.config import VllmConfig

10 from vllm.entrypoints.openai.protocol import IOProcessorResponse

11 from vllm.plugins.io_processors import get_io_processor

12 -from vllm.pooling_params import PoolingParams

12

13 MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"

14

@@ -94,8 +93,6 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):

93 out_data_format="b64_json",

94 )

95

97 - pooling_params = PoolingParams(activation=False)

98 -

96 with vllm_runner(

97 model_name,

98 runner="pooling",

@@ -109,9 +106,7 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):

106 model_impl="terratorch",

107 io_processor_plugin="prithvi_to_tiff",

108 ) as llm_runner:

112 - pooler_output = llm_runner.get_llm().encode(

113 - img_prompt, pooling_params=pooling_params, pooling_task="token_classify"

114 - )

109 + pooler_output = llm_runner.get_llm().encode(img_prompt, pooling_task="plugin")

110 output = pooler_output[0].outputs

111

112 # verify the output is formatted as expected for this plugin

vllm/entrypoints/llm.py

+28 -13

@@ -1024,19 +1024,6 @@ class LLM:

1024 "pooling model."

1025 )

1026

1027 - if pooling_task not in self.supported_tasks:

1028 - raise ValueError(f"pooling_task must be one of {self.supported_tasks}.")

1029 -

1030 - if pooling_params is None:

1031 - # Use default pooling params.

1032 - pooling_params = PoolingParams()

1033 -

1034 - for param in as_iter(pooling_params):

1035 - param.verify(pooling_task, model_config)

1036 - # for backwards compatibility

1037 - if truncate_prompt_tokens is not None:

1038 - param.truncate_prompt_tokens = truncate_prompt_tokens

1039 -

1027 io_processor_prompt = False

1028 if isinstance(prompts, dict) and "data" in prompts:

1029 io_processor_prompt = True

@@ -1054,6 +1041,34 @@ class LLM:

1041 # obtain the actual model prompts from the pre-processor

1042 prompts = self.io_processor.pre_process(prompt=validated_prompt)

1043

1044 + if io_processor_prompt:

1045 + assert self.io_processor is not None

1046 + if is_list_of(pooling_params, PoolingParams):

1047 + validated_pooling_params: list[PoolingParams] = []

1048 + for param in as_iter(pooling_params):

1049 + validated_pooling_params.append(

1050 + self.io_processor.validate_or_generate_params(param)

1051 + )

1052 + pooling_params = validated_pooling_params

1053 + else:

1054 + assert not isinstance(pooling_params, Sequence)

1055 + pooling_params = self.io_processor.validate_or_generate_params(

1056 + pooling_params

1057 + )

1058 + else:

1059 + if pooling_params is None:

1060 + # Use default pooling params.

1061 + pooling_params = PoolingParams()

1062 +

1063 + if pooling_task not in self.supported_tasks:

1064 + raise ValueError(f"pooling_task must be one of {self.supported_tasks}.")

1065 +

1066 + for param in as_iter(pooling_params):

1067 + param.verify(pooling_task, model_config)

1068 + # for backwards compatibility

1069 + if truncate_prompt_tokens is not None:

1070 + param.truncate_prompt_tokens = truncate_prompt_tokens

1071 +

1072 self._validate_and_add_requests(

1073 prompts=prompts,

1074 params=pooling_params,

vllm/entrypoints/openai/api_server.py

+6 -1

@@ -1748,7 +1748,12 @@ async def init_app_state(

1748 log_error_stack=args.log_error_stack,

1749 )

1750 )

1751 - if ("token_embed" in supported_tasks or "token_classify" in supported_tasks)

1751 + if (

1752 + any(

1753 + task in supported_tasks

1754 + for task in ["token_embed", "token_classify", "plugin"]

1755 + )

1756 + )

1757 else None

1758 )

1759 state.openai_serving_embedding = (

vllm/entrypoints/openai/protocol.py

+1 -6

@@ -1707,11 +1707,6 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):

1707 if the served model does not use priority scheduling.

1708 """

1709 data: T

1710 - """

1711 - When using plugins IOProcessor plugins, the actual input is processed

1712 - by the plugin itself. Hence, we use a generic type for the request data

1713 - """

1714 - activation: bool = False

1710

1711 encoding_format: EncodingFormat = "float"

1712 embed_dtype: EmbedDType = Field(

@@ -1732,7 +1727,7 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):

1727 )

1728

1729 def to_pooling_params(self):

1735 - return PoolingParams(task="token_classify", activation=self.activation)

1730 + return PoolingParams()

1731

1732

1733 class IOProcessorResponse(OpenAIBaseModel, Generic[T]):

vllm/entrypoints/openai/serving_pooling.py

+11 -2

@@ -32,7 +32,7 @@ from vllm.entrypoints.renderer import RenderConfig

32 from vllm.entrypoints.utils import _validate_truncation_size

33 from vllm.logger import init_logger

34 from vllm.outputs import PoolingRequestOutput

35 -from vllm.tasks import SupportedTask

35 +from vllm.tasks import PoolingTask, SupportedTask

36 from vllm.utils.async_utils import merge_async_iterators

37 from vllm.utils.serial_utils import (

38 EmbedDType,

@@ -161,12 +161,21 @@ class OpenAIServingPooling(OpenAIServing):

161 # Schedule the request and get the result generator.

162 generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

163 try:

164 - pooling_params = request.to_pooling_params()

164 + if is_io_processor_request:

165 + assert self.io_processor is not None and isinstance(

166 + request, IOProcessorRequest

167 + )

168 + pooling_params = self.io_processor.validate_or_generate_params()

169 + else:

170 + pooling_params = request.to_pooling_params()

171

172 + pooling_task: PoolingTask

173 if "token_embed" in self.supported_tasks:

174 pooling_task = "token_embed"

175 elif "token_classify" in self.supported_tasks:

176 pooling_task = "token_classify"

177 + elif "plugin" in self.supported_tasks:

178 + pooling_task = "plugin"

179 else:

180 return self.create_error_response(

181 f"pooling_task must be one of {self.supported_tasks}."

vllm/model_executor/layers/pooler.py

+12

@@ -414,6 +414,18 @@ class Pooler(nn.Module, ABC):

414 raise NotImplementedError

415

416

417 +class DummyPooler(Pooler):

418 + def get_supported_tasks(self) -> Set[PoolingTask]:

419 + return {"plugin", "score"}

420 +

421 + def forward(

422 + self,

423 + hidden_states: list[torch.Tensor] | torch.Tensor,

424 + pooling_metadata: PoolingMetadata,

425 + ) -> PoolerOutput:

426 + return hidden_states

427 +

428 +

429 class PoolerHead(nn.Module):

430 def __init__(self, activation: PoolerActivation) -> None:

431 super().__init__()

vllm/model_executor/models/terratorch.py

+2 -4

@@ -34,7 +34,7 @@ from transformers import BatchFeature

34 from vllm.config import VllmConfig

35 from vllm.config.multimodal import BaseDummyOptions

36 from vllm.logger import init_logger

37 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler

37 +from vllm.model_executor.layers.pooler import DispatchPooler, DummyPooler

38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

39 from vllm.model_executor.models.utils import AutoWeightsLoader

40 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -249,9 +249,7 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

249 pooler_config = vllm_config.model_config.pooler_config

250 assert pooler_config is not None

251

252 - self.pooler = DispatchPooler(

253 - {"token_classify": Pooler.for_token_classify(pooler_config)}

254 - )

252 + self.pooler = DispatchPooler({"plugin": DummyPooler()})

253

254 def get_input_embeddings(

255 self,

vllm/plugins/io_processors/interface.py

+7

@@ -9,6 +9,8 @@ from vllm.config import VllmConfig

9 from vllm.entrypoints.openai.protocol import IOProcessorResponse

10 from vllm.inputs.data import PromptType

11 from vllm.outputs import PoolingRequestOutput

12 +from vllm.pooling_params import PoolingParams

13 +from vllm.sampling_params import SamplingParams

14

15 IOProcessorInput = TypeVar("IOProcessorInput")

16 IOProcessorOutput = TypeVar("IOProcessorOutput")

@@ -63,6 +65,11 @@ class IOProcessor(ABC, Generic[IOProcessorInput, IOProcessorOutput]):

65 def parse_request(self, request: Any) -> IOProcessorInput:

66 raise NotImplementedError

67

68 + def validate_or_generate_params(

69 + self, params: SamplingParams | PoolingParams | None = None

70 + ) -> SamplingParams | PoolingParams:

71 + return params or PoolingParams()

72 +

73 @abstractmethod

74 def output_to_response(

75 self, plugin_output: IOProcessorOutput

vllm/pooling_params.py

+5

@@ -84,6 +84,11 @@ class PoolingParams(

84 msg = f"You cannot overwrite {self.task=!r} with {task=!r}!"

85 raise ValueError(msg)

86

87 + # plugin task uses io_processor.parse_request to verify inputs,

88 + # skipping PoolingParams verify

89 + if self.task == "plugin":

90 + return

91 +

92 # NOTE: Task validation needs to done against the model instance,

93 # which is not available in model config. So, it's not included

94 # in this method

vllm/tasks.py

+3 -1

@@ -5,7 +5,9 @@ from typing import Literal, get_args

5 GenerationTask = Literal["generate", "transcription"]

6 GENERATION_TASKS = get_args(GenerationTask)

7

8 -PoolingTask = Literal["embed", "classify", "score", "token_embed", "token_classify"]

8 +PoolingTask = Literal[

9 + "embed", "classify", "score", "token_embed", "token_classify", "plugin"

10 +]

11 POOLING_TASKS = get_args(PoolingTask)

12

13 SupportedTask = Literal[GenerationTask, PoolingTask]

[Frontend][4/N] Improve all pooling task | Add plugin pooling task (#26973)