@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 45badd05d0

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 01:00 UTC

[Core] Set pooling params based on task and model (#21128)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jul 18, 2025 at 20:41 UTC 45badd05d04254eb75c48cee7b1d454a80de2165

24 files changed +498 -230

tests/models/language/pooling/test_gritlm.py

+11 -15

@@ -2,9 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from __future__ import annotations

5 -import importlib.util

6 -from array import array

7 -

5 +import numpy as np

6 import openai

7 import pytest

8 from scipy.spatial.distance import cosine

@@ -14,10 +12,6 @@ from vllm.config import ModelConfig

13 from ....utils import RemoteOpenAIServer

17 -# GritLM embedding implementation is only supported by XFormers backend.

18 -pytestmark = pytest.mark.skipif(not importlib.util.find_spec("xformers"),

19 - reason="GritLM requires XFormers")

20 -

15 MODEL_NAME = "parasail-ai/GritLM-7B-vllm"

16 MAX_MODEL_LEN = 4000

@@ -26,11 +20,11 @@ def _arr(arr):

20 """

21 Convert a list of integers to an array of integers.

22 """

29 - return array("i", arr)

23 + return np.array(arr)

26 def test_find_array():

33 - from vllm.model_executor.models.gritlm import GritLMPooler

27 + from vllm.model_executor.models.gritlm import GritLMMeanPool

29 model_config = ModelConfig(

30 MODEL_NAME,

@@ -41,17 +35,19 @@ def test_find_array():

35 dtype="bfloat16",

36 seed=0,

37 )

44 - pooler = GritLMPooler(model_config=model_config)

38 + pooling = GritLMMeanPool(model_config=model_config)

40 arr = _arr([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

48 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3

49 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3

50 - assert pooler._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1

51 - assert pooler._find_array(arr, _arr([3, 5]), start_idx=0) == -1

42 + assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=0) == 3

43 + assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=1) == 3

44 + assert pooling._find_array(arr, _arr([3, 4, 5]), start_idx=5) == -1

45 + assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=3) == -1

46 + assert pooling._find_array(arr, _arr([3, 4, 5]), end_idx=4) == 3

47 + assert pooling._find_array(arr, _arr([3, 5]), start_idx=0) == -1

49 with pytest.raises(ValueError):

54 - pooler._find_array(arr, _arr([3, 4, 5]), start_idx=-1)

50 + pooling._find_array(arr, _arr([3, 4, 5]), start_idx=-1)

53 def run_llm_encode(

vllm/entrypoints/llm.py

+33 -16

@@ -44,7 +44,7 @@ from vllm.model_executor.layers.quantization import QuantizationMethods

44 from vllm.outputs import (ClassificationRequestOutput, EmbeddingRequestOutput,

45 PoolingRequestOutput, RequestOutput,

46 ScoringRequestOutput)

47 -from vllm.pooling_params import PoolingParams

47 +from vllm.pooling_params import PoolingParams, PoolingTask

48 from vllm.prompt_adapter.request import PromptAdapterRequest

49 from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams,

50 RequestOutputKind, SamplingParams)

@@ -964,6 +964,7 @@ class LLM:

964 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

965 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

966 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

967 + pooling_task: PoolingTask = "encode",

968 ) -> list[PoolingRequestOutput]:

969 ...

970

@@ -979,6 +980,7 @@ class LLM:

980 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

981 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

982 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

983 + pooling_task: PoolingTask = "encode",

984 ) -> list[PoolingRequestOutput]:

985 ...

986

@@ -994,6 +996,7 @@ class LLM:

996 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

997 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

998 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

999 + pooling_task: PoolingTask = "encode",

1000 ) -> list[PoolingRequestOutput]:

1001 ...

1002

@@ -1010,6 +1013,7 @@ class LLM:

1013 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1014 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1015 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1016 + pooling_task: PoolingTask = "encode",

1017 ) -> list[PoolingRequestOutput]:

1018 ...

1019

@@ -1026,6 +1030,7 @@ class LLM:

1030 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1031 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1032 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1033 + pooling_task: PoolingTask = "encode",

1034 ) -> list[PoolingRequestOutput]:

1035 ...

1036

@@ -1040,6 +1045,7 @@ class LLM:

1045 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1046 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1047 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1048 + pooling_task: PoolingTask = "encode",

1049 ) -> list[PoolingRequestOutput]:

1050 ...

1051

@@ -1059,6 +1065,7 @@ class LLM:

1065 use_tqdm: Union[bool, Callable[..., tqdm]] = True,

1066 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

1067 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

1068 + pooling_task: PoolingTask = "encode",

1069 ) -> list[PoolingRequestOutput]:

1070 """Apply pooling to the hidden states corresponding to the input

1071 prompts.

@@ -1080,6 +1087,7 @@ class LLM:

1087 lora_request: LoRA request to use for generation, if any.

1088 prompt_adapter_request: Prompt Adapter request to use for

1089 generation, if any.

1090 + pooling_task: Override the pooling task to use.

1091

1092 Returns:

1093 A list of `PoolingRequestOutput` objects containing the

@@ -1116,11 +1124,12 @@ class LLM:

1124 if pooling_params is None:

1125 # Use default pooling params.

1126 pooling_params = PoolingParams()

1119 - elif isinstance(pooling_params, PoolingParams):

1120 - pooling_params.verify(model_config)

1127 +

1128 + if isinstance(pooling_params, PoolingParams):

1129 + pooling_params.verify(pooling_task, model_config)

1130 else:

1131 for pooling_param in pooling_params:

1123 - pooling_param.verify(model_config)

1132 + pooling_param.verify(pooling_task, model_config)

1133

1134 tokenization_kwargs = dict[str, Any]()

1135 _validate_truncation_size(model_config.max_model_len,

@@ -1181,12 +1190,15 @@ class LLM:

1190 raise ValueError("Embedding API is not supported by this model. "

1191 "Please set `--task embed`.")

1192

1184 - items = self.encode(prompts,

1185 - truncate_prompt_tokens=truncate_prompt_tokens,

1186 - use_tqdm=use_tqdm,

1187 - pooling_params=pooling_params,

1188 - lora_request=lora_request,

1189 - prompt_adapter_request=prompt_adapter_request)

1193 + items = self.encode(

1194 + prompts,

1195 + truncate_prompt_tokens=truncate_prompt_tokens,

1196 + use_tqdm=use_tqdm,

1197 + pooling_params=pooling_params,

1198 + lora_request=lora_request,

1199 + prompt_adapter_request=prompt_adapter_request,

1200 + pooling_task="embed",

1201 + )

1202

1203 return [EmbeddingRequestOutput.from_base(item) for item in items]

1204

@@ -1228,10 +1240,13 @@ class LLM:

1240 "Classification API is not supported by this model. "

1241 "Please set `--task classify`.")

1242

1231 - items = self.encode(prompts,

1232 - use_tqdm=use_tqdm,

1233 - lora_request=lora_request,

1234 - prompt_adapter_request=prompt_adapter_request)

1243 + items = self.encode(

1244 + prompts,

1245 + use_tqdm=use_tqdm,

1246 + lora_request=lora_request,

1247 + prompt_adapter_request=prompt_adapter_request,

1248 + pooling_task="classify",

1249 + )

1250

1251 return [ClassificationRequestOutput.from_base(item) for item in items]

1252

@@ -1251,7 +1266,9 @@ class LLM:

1266 truncate_prompt_tokens=truncate_prompt_tokens,

1267 use_tqdm=use_tqdm,

1268 lora_request=lora_request,

1254 - prompt_adapter_request=prompt_adapter_request)

1269 + prompt_adapter_request=prompt_adapter_request,

1270 + pooling_task="embed",

1271 + )

1272

1273 encoded_output_1: list[PoolingRequestOutput] = encoded_output[

1274 0:len(text_1)]

@@ -1287,7 +1304,7 @@ class LLM:

1304 if len(data_1) == 1:

1305 data_1 = data_1 * len(data_2)

1306

1290 - pooling_params = PoolingParams(use_cross_encoder=True)

1307 + pooling_params = PoolingParams(task="score")

1308 tokenization_kwargs: dict[str, Any] = {}

1309 _validate_truncation_size(self.llm_engine.model_config.max_model_len,

1310 truncate_prompt_tokens, tokenization_kwargs)

vllm/entrypoints/openai/protocol.py

+4 -4

@@ -1347,8 +1347,8 @@ class ScoreRequest(OpenAIBaseModel):

1347

1348 # --8<-- [end:score-extra-params]

1349

1350 - def to_pooling_params(self, *, use_cross_encoder: bool = False):

1351 - return PoolingParams(use_cross_encoder=use_cross_encoder)

1350 + def to_pooling_params(self):

1351 + return PoolingParams()

1352

1353

1354 class RerankRequest(OpenAIBaseModel):

@@ -1375,8 +1375,8 @@ class RerankRequest(OpenAIBaseModel):

1375

1376 # --8<-- [end:rerank-extra-params]

1377

1378 - def to_pooling_params(self, *, use_cross_encoder: bool = False):

1379 - return PoolingParams(use_cross_encoder=use_cross_encoder)

1378 + def to_pooling_params(self):

1379 + return PoolingParams()

1380

1381

1382 class RerankDocument(BaseModel):

vllm/entrypoints/openai/serving_classification.py

+32

@@ -6,6 +6,7 @@ from typing import Optional, Union, cast

7 import numpy as np

8 from fastapi import Request

9 +from typing_extensions import override

11 from vllm.config import ModelConfig

12 from vllm.engine.protocol import EngineClient

@@ -21,12 +22,14 @@ from vllm.entrypoints.openai.serving_engine import (ClassificationServeContext,

22 from vllm.entrypoints.openai.serving_models import OpenAIServingModels

23 from vllm.logger import init_logger

24 from vllm.outputs import ClassificationOutput, PoolingRequestOutput

25 +from vllm.pooling_params import PoolingParams

27 logger = init_logger(__name__)

30 class ClassificationMixin(OpenAIServing):

32 + @override

33 async def _preprocess(

34 self,

35 ctx: ServeContext,

@@ -75,6 +78,7 @@ class ClassificationMixin(OpenAIServing):

78 logger.exception("Error in preprocessing prompt inputs")

79 return self.create_error_response(str(e))

81 + @override

82 def _build_response(

83 self,

84 ctx: ServeContext,

@@ -158,3 +162,31 @@ class ServingClassification(ClassificationMixin):

162 )

163

164 return await super().handle(ctx) # type: ignore

165 +

166 + @override

167 + def _validate_request(

168 + self,

169 + ctx: ClassificationServeContext,

170 + ) -> Optional[ErrorResponse]:

171 + if error := super()._validate_request(ctx):

172 + return error

173 +

174 + ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens

175 +

176 + return None

177 +

178 + @override

179 + def _create_pooling_params(

180 + self,

181 + ctx: ClassificationServeContext,

182 + ) -> Union[PoolingParams, ErrorResponse]:

183 + pooling_params = super()._create_pooling_params(ctx)

184 + if isinstance(pooling_params, ErrorResponse):

185 + return pooling_params

186 +

187 + try:

188 + pooling_params.verify("classify", self.model_config)

189 + except ValueError as e:

190 + return self.create_error_response(str(e))

191 +

192 + return pooling_params

vllm/entrypoints/openai/serving_embedding.py

+15 -3

@@ -24,6 +24,7 @@ from vllm.entrypoints.openai.serving_models import OpenAIServingModels

24 from vllm.logger import init_logger

25 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,

26 PoolingRequestOutput)

27 +from vllm.pooling_params import PoolingParams

29 logger = init_logger(__name__)

@@ -45,6 +46,7 @@ def _get_embedding(

47 class EmbeddingMixin(OpenAIServing):

49 + @override

50 async def _preprocess(

51 self,

52 ctx: ServeContext,

@@ -97,6 +99,7 @@ class EmbeddingMixin(OpenAIServing):

99 logger.exception("Error in preprocessing prompt inputs")

100 return self.create_error_response(str(e))

101

102 + @override

103 def _build_response(

104 self,

105 ctx: ServeContext,

@@ -191,11 +194,20 @@ class OpenAIServingEmbedding(EmbeddingMixin):

194

195 ctx.truncate_prompt_tokens = ctx.request.truncate_prompt_tokens

196

194 - pooling_params = ctx.request.to_pooling_params()

197 + return None

198 +

199 + @override

200 + def _create_pooling_params(

201 + self,

202 + ctx: ServeContext[EmbeddingRequest],

203 + ) -> Union[PoolingParams, ErrorResponse]:

204 + pooling_params = super()._create_pooling_params(ctx)

205 + if isinstance(pooling_params, ErrorResponse):

206 + return pooling_params

207

208 try:

197 - pooling_params.verify(self.model_config)

209 + pooling_params.verify("embed", self.model_config)

210 except ValueError as e:

211 return self.create_error_response(str(e))

212

201 - return None

213 + return pooling_params

vllm/entrypoints/openai/serving_engine.py

+13 -5

@@ -305,6 +305,16 @@ class OpenAIServing:

305 " Please, select a smaller truncation size.")

306 return None

307

308 + def _create_pooling_params(

309 + self,

310 + ctx: ServeContext,

311 + ) -> Union[PoolingParams, ErrorResponse]:

312 + if not hasattr(ctx.request, "to_pooling_params"):

313 + return self.create_error_response(

314 + "Request type does not support pooling parameters")

315 +

316 + return ctx.request.to_pooling_params()

317 +

318 async def _prepare_generators(

319 self,

320 ctx: ServeContext,

@@ -318,11 +328,9 @@ class OpenAIServing:

328 trace_headers = (None if ctx.raw_request is None else await

329 self._get_trace_headers(ctx.raw_request.headers))

330

321 - if not hasattr(ctx.request, "to_pooling_params"):

322 - return self.create_error_response(

323 - "Request type does not support pooling parameters")

324 -

325 - pooling_params = ctx.request.to_pooling_params()

331 + pooling_params = self._create_pooling_params(ctx)

332 + if isinstance(pooling_params, ErrorResponse):

333 + return pooling_params

334

335 if ctx.engine_prompts is None:

336 return self.create_error_response(

vllm/entrypoints/openai/serving_pooling.py

@@ -142,6 +142,11 @@ class OpenAIServingPooling(OpenAIServing):

142 try:

143 pooling_params = request.to_pooling_params()

144

145 + try:

146 + pooling_params.verify("encode", self.model_config)

147 + except ValueError as e:

148 + return self.create_error_response(str(e))

149 +

150 for i, engine_prompt in enumerate(engine_prompts):

151 request_id_item = f"{request_id}-{i}"

152

vllm/entrypoints/openai/serving_score.py

+21 -9

@@ -55,14 +55,13 @@ class ServingScores(OpenAIServing):

55 texts_1: list[str],

56 texts_2: list[str],

57 request: Union[RerankRequest, ScoreRequest],

58 - request_id=str,

58 + request_id: str,

59 tokenization_kwargs: Optional[dict[str, Any]] = None,

60 lora_request: Optional[Union[LoRARequest, None]] = None,

61 prompt_adapter_request: Optional[Union[PromptAdapterRequest,

62 None]] = None,

63 trace_headers: Optional[Mapping[str, str]] = None,

64 - ) -> list[PoolingRequestOutput]:

65 -

64 + ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

65 input_texts = texts_1 + texts_2

67 engine_prompts: list[TokensPrompt] = []

@@ -89,6 +88,11 @@ class ServingScores(OpenAIServing):

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

89 pooling_params = request.to_pooling_params()

91 + try:

92 + pooling_params.verify("embed", self.model_config)

93 + except ValueError as e:

94 + return self.create_error_response(str(e))

95 +

96 for i, engine_prompt in enumerate(engine_prompts):

98 request_id_item = f"{request_id}-{i}"

@@ -169,14 +173,13 @@ class ServingScores(OpenAIServing):

173 data_1: Union[list[str], list[ScoreContentPartParam]],

174 data_2: Union[list[str], list[ScoreContentPartParam]],

175 request: Union[RerankRequest, ScoreRequest],

172 - request_id=str,

176 + request_id: str,

177 tokenization_kwargs: Optional[dict[str, Any]] = None,

178 lora_request: Optional[Union[LoRARequest, None]] = None,

179 prompt_adapter_request: Optional[Union[PromptAdapterRequest,

180 None]] = None,

181 trace_headers: Optional[Mapping[str, str]] = None,

178 - ) -> list[PoolingRequestOutput]:

179 -

182 + ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

183 request_prompts: list[str] = []

184 engine_prompts: list[TokensPrompt] = []

185

@@ -245,7 +248,12 @@ class ServingScores(OpenAIServing):

248 # Schedule the request and get the result generator.

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

250

248 - pooling_params = request.to_pooling_params(use_cross_encoder=True)

251 + pooling_params = request.to_pooling_params()

252 +

253 + try:

254 + pooling_params.verify("score", self.model_config)

255 + except ValueError as e:

256 + return self.create_error_response(str(e))

257

258 for i, engine_prompt in enumerate(engine_prompts):

259 request_id_item = f"{request_id}-{i}"

@@ -286,8 +294,7 @@ class ServingScores(OpenAIServing):

294 request_id: str,

295 raw_request: Optional[Request] = None,

296 truncate_prompt_tokens: Optional[int] = None,

289 - ) -> list[PoolingRequestOutput]:

290 -

297 + ) -> Union[list[PoolingRequestOutput], ErrorResponse]:

298 (

299 lora_request,

300 prompt_adapter_request,

@@ -374,6 +381,8 @@ class ServingScores(OpenAIServing):

381 raw_request,

382 request.truncate_prompt_tokens,

383 )

384 + if isinstance(final_res_batch, ErrorResponse):

385 + return final_res_batch

386

387 return self.request_output_to_score_response(

388 final_res_batch,

@@ -420,6 +429,9 @@ class ServingScores(OpenAIServing):

429 raw_request,

430 request.truncate_prompt_tokens,

431 )

432 + if isinstance(final_res_batch, ErrorResponse):

433 + return final_res_batch

434 +

435 return self.request_output_to_rerank_response(

436 final_res_batch,

437 request_id,

vllm/executor/executor_base.py

@@ -4,6 +4,7 @@

4 import asyncio

5 import time

6 from abc import ABC, abstractmethod

7 +from functools import cached_property

8 from typing import (Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple,

9 Union)

@@ -15,6 +16,7 @@ from vllm.config import VllmConfig

16 from vllm.logger import init_logger

17 from vllm.lora.request import LoRARequest

18 from vllm.model_executor.layers.sampler import SamplerOutput

19 +from vllm.pooling_params import PoolingTask

20 from vllm.prompt_adapter.request import PromptAdapterRequest

21 from vllm.sequence import ExecuteModelRequest, PoolerOutput

22 from vllm.utils import make_async

@@ -135,6 +137,11 @@ class ExecutorBase(ABC):

137

138 return self.collective_rpc(rpc_func)

139

140 + @cached_property # Avoid unnecessary RPC calls

141 + def supported_pooling_tasks(self) -> tuple[PoolingTask, ...]:

142 + output = self.collective_rpc("get_supported_pooling_tasks")

143 + return tuple({task for tasks in output for task in tasks})

144 +

145 def execute_model(

146 self, execute_model_req: ExecuteModelRequest

147 ) -> Optional[List[Union[SamplerOutput, PoolerOutput]]]:

vllm/model_executor/layers/pooler.py

+95 -54

@@ -3,7 +3,7 @@

3 from abc import ABC, abstractmethod

4 from dataclasses import dataclass

5 from enum import IntEnum

6 -from typing import Callable, Literal, Optional, TypeVar, Union

6 +from typing import Callable, Optional, TypeVar, Union

8 import torch

9 import torch.nn as nn

@@ -15,13 +15,12 @@ from vllm.config import ModelConfig, PoolerConfig

15 from vllm.model_executor.pooling_metadata import ( # noqa: E501

16 PoolingMetadata as V0PoolingMetadata)

17 from vllm.model_executor.pooling_metadata import PoolingTensors

18 -from vllm.pooling_params import PoolingParams

18 +from vllm.pooling_params import PoolingParams, PoolingTask

19 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput

20 from vllm.utils import resolve_obj_by_qualname

21 from vllm.v1.pool.metadata import PoolingMetadata as V1PoolingMetadata

23 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]

24 -PoolingTask = Literal["encode", "embed", "classify", "score"]

26 class PoolingType(IntEnum):

@@ -67,6 +66,15 @@ class ResolvedPoolingConfig:

66 )

69 +@dataclass(frozen=True)

70 +class PoolingParamsUpdate:

71 + requires_token_ids: bool = False

72 + """Set this flag to enable `get_prompt_token_ids` for your pooler."""

73 +

74 + def apply(self, params: PoolingParams) -> None:

75 + params.requires_token_ids = self.requires_token_ids

76 +

77 +

78 class Pooler(nn.Module, ABC):

79 """The interface required for all poolers used in pooling models in vLLM."""

@@ -93,7 +101,10 @@ class Pooler(nn.Module, ABC):

101

102 return SimplePooler.from_config(resolved_config)

103

96 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

104 + def get_pooling_updates(

105 + self,

106 + task: PoolingTask,

107 + ) -> Optional[PoolingParamsUpdate]:

108 """

109 Construct the pooling parameters to use for a task,

110 or `None` if the task is not supported.

@@ -121,6 +132,23 @@ def get_prompt_lens(

132 pooling_metadata, hidden_states.device).prompt_lens

133

134

135 +def get_prompt_token_ids(

136 + pooling_metadata: PoolingMetadata) -> list[torch.Tensor]:

137 + if isinstance(pooling_metadata, V1PoolingMetadata):

138 + assert pooling_metadata.prompt_token_ids is not None, (

139 + "Please set `requires_token_ids=True` in `get_pooling_updates`")

140 +

141 + return [

142 + pooling_metadata.prompt_token_ids[i, :num]

143 + for i, num in enumerate(pooling_metadata.prompt_lens)

144 + ]

145 +

146 + return [

147 + torch.tensor(seq_data_i.prompt_token_ids)

148 + for seq_data_i in pooling_metadata.seq_data.values()

149 + ]

150 +

151 +

152 def get_classification_activation_function(config: PretrainedConfig):

153 return PoolerClassify()

154

@@ -165,7 +193,10 @@ class PoolingMethod(nn.Module, ABC):

193 raise NotImplementedError(f"Unsupported method: {pooling_type}")

194

195 @abstractmethod

168 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

196 + def get_pooling_updates(

197 + self,

198 + task: PoolingTask,

199 + ) -> Optional[PoolingParamsUpdate]:

200 raise NotImplementedError

201

202 @abstractmethod

@@ -206,11 +237,14 @@ class PoolingMethod(nn.Module, ABC):

237

238 class CLSPool(PoolingMethod):

239

209 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

240 + def get_pooling_updates(

241 + self,

242 + task: PoolingTask,

243 + ) -> Optional[PoolingParamsUpdate]:

244 # The equalities are split up to keep mypy happy

245 if (task == "encode" or task == "embed" or task == "classify"

246 or task == "score"):

213 - return PoolingParams()

247 + return PoolingParamsUpdate()

248

249 assert_never(task)

250

@@ -236,11 +270,14 @@ class CLSPool(PoolingMethod):

270

271 class LastPool(PoolingMethod):

272

239 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

273 + def get_pooling_updates(

274 + self,

275 + task: PoolingTask,

276 + ) -> Optional[PoolingParamsUpdate]:

277 # The equalities are split up to keep mypy happy

278 if (task == "encode" or task == "embed" or task == "classify"

279 or task == "score"):

243 - return PoolingParams()

280 + return PoolingParamsUpdate()

281

282 assert_never(task)

283

@@ -262,9 +299,12 @@ class LastPool(PoolingMethod):

299

300 class AllPool(PoolingMethod):

301

265 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

302 + def get_pooling_updates(

303 + self,

304 + task: PoolingTask,

305 + ) -> Optional[PoolingParamsUpdate]:

306 if task == "encode":

267 - return PoolingParams()

307 + return PoolingParamsUpdate()

308

309 # The equalities are split up to keep mypy happy

310 if task == "embed" or task == "classify" or task == "score":

@@ -299,11 +339,14 @@ class AllPool(PoolingMethod):

339

340 class MeanPool(PoolingMethod):

341

302 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

342 + def get_pooling_updates(

343 + self,

344 + task: PoolingTask,

345 + ) -> Optional[PoolingParamsUpdate]:

346 # The equalities are split up to keep mypy happy

347 if (task == "encode" or task == "embed" or task == "classify"

348 or task == "score"):

306 - return PoolingParams()

349 + return PoolingParamsUpdate()

350

351 assert_never(task)

352

@@ -520,8 +563,11 @@ class SimplePooler(Pooler):

563 self.pooling = pooling

564 self.head = head

565

523 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

524 - return self.pooling.get_pooling_params(task)

566 + def get_pooling_updates(

567 + self,

568 + task: PoolingTask,

569 + ) -> Optional[PoolingParamsUpdate]:

570 + return self.pooling.get_pooling_updates(task)

571

572 def forward(

573 self,

@@ -559,27 +605,13 @@ class StepPooler(Pooler):

605 self.step_tag_id = step_tag_id

606 self.returned_token_ids = returned_token_ids

607

562 - def get_prompt_token_ids(

563 - self,

564 - pooling_metadata: PoolingMetadata,

565 - ) -> list[torch.Tensor]:

566 - if isinstance(pooling_metadata, V1PoolingMetadata):

567 - return [

568 - pooling_metadata.prompt_token_ids[i, :num]

569 - for i, num in enumerate(pooling_metadata.prompt_lens)

570 - ]

571 - return [

572 - torch.tensor(seq_data_i.prompt_token_ids)

573 - for seq_data_i in pooling_metadata.seq_data.values()

574 - ]

575 -

608 def extract_states(

609 self,

610 hidden_states: Union[torch.Tensor, list[torch.Tensor]],

611 pooling_metadata: PoolingMetadata,

612 ) -> Union[list[torch.Tensor], torch.Tensor]:

613 pooled_data_lst = self.pooling(hidden_states, pooling_metadata)

582 - prompt_token_ids = self.get_prompt_token_ids(pooling_metadata)

614 + prompt_token_ids = get_prompt_token_ids(pooling_metadata)

615

616 pooled_data = list[torch.Tensor]()

617 returned_token_ids = self.returned_token_ids

@@ -595,9 +627,12 @@ class StepPooler(Pooler):

627

628 return pooled_data

629

598 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

630 + def get_pooling_updates(

631 + self,

632 + task: PoolingTask,

633 + ) -> Optional[PoolingParamsUpdate]:

634 if task == "encode":

600 - return PoolingParams(logits_processing_needs_token_ids=True)

635 + return PoolingParamsUpdate(requires_token_ids=True)

636

637 # The equalities are split up to keep mypy happy

638 if task == "embed" or task == "classify" or task == "score":

@@ -650,19 +685,24 @@ class ClassifierPooler(nn.Module):

685 self.cross_encoder_act_fn = get_cross_encoder_activation_function(

686 config.hf_config) if act_fn is None else act_fn

687

653 - def _get_act_fn(self, use_cross_encoder: bool):

654 - return (self.cross_encoder_act_fn

655 - if use_cross_encoder else self.classification_act_fn)

688 + def _get_act_fn(self, task: PoolingTask):

689 + if task == "encode" or task == "classify":

690 + return self.classification_act_fn

691 + if task == "score":

692 + return self.cross_encoder_act_fn

693 +

694 + raise ValueError(f"Unsupported task: {task!r}")

695 +

696 + def get_pooling_updates(

697 + self,

698 + task: PoolingTask,

699 + ) -> Optional[PoolingParamsUpdate]:

700 + # The equalities are split up to keep mypy happy

701 + if task == "encode" or task == "classify" or task == "score":

702 + return PoolingParamsUpdate()

703

657 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

658 - if task == "encode":

659 - return PoolingParams()

704 if task == "embed":

705 return None

662 - if task == "classify":

663 - return PoolingParams()

664 - if task == "score":

665 - return PoolingParams(use_cross_encoder=True)

706

707 assert_never(task)

708

@@ -682,27 +722,28 @@ class ClassifierPooler(nn.Module):

722 else:

723 pooled_output = [self.classifier(data) for data in pooled_data]

724

725 + task_list: list[PoolingTask]

726 if isinstance(pooling_metadata, V0PoolingMetadata):

686 - use_cross_encoder_list = [

687 - pooling_param.use_cross_encoder

688 - for _, pooling_param in pooling_metadata.seq_groups

727 + task_list = [

728 + task for _, pooling_param in pooling_metadata.seq_groups

729 + if (task := pooling_param.task) is not None

730 ]

731 else:

691 - use_cross_encoder_list = [

692 - pooling_param.use_cross_encoder

693 - for pooling_param in pooling_metadata.pooling_params

732 + task_list = [

733 + task for pooling_param in pooling_metadata.pooling_params

734 + if (task := pooling_param.task) is not None

735 ]

736

737 + assert len(task_list) == len(pooled_output)

738 +

739 # shape of scores: (batch_size, num_labels)

697 - if all(use_cross_encoder == use_cross_encoder_list[0]

698 - for use_cross_encoder in use_cross_encoder_list):

699 - act_fn = self._get_act_fn(use_cross_encoder_list[0])

740 + if len(set(task_list)) <= 1:

741 + act_fn = self._get_act_fn(task_list[0])

742 scores = act_fn(pooled_output)

743 else:

744 scores = torch.stack([

703 - self._get_act_fn(use_cross_encoder)(vecs)

704 - for use_cross_encoder, vecs in zip(use_cross_encoder_list,

705 - pooled_output)

745 + self._get_act_fn(task)(vecs)

746 + for task, vecs in zip(task_list, pooled_output)

747 ])

748

749 return build_output(scores)

vllm/model_executor/models/bert.py

+8 -4

@@ -18,13 +18,14 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

18 QKVParallelLinear,

19 RowParallelLinear)

20 from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,

21 - PoolingMethod, PoolingTask,

21 + PoolingMethod,

22 + PoolingParamsUpdate,

23 PoolingType)

24 from vllm.model_executor.layers.quantization import QuantizationConfig

25 from vllm.model_executor.layers.vocab_parallel_embedding import (

26 VocabParallelEmbedding)

27 from vllm.model_executor.pooling_metadata import PoolingMetadata

27 -from vllm.pooling_params import PoolingParams

28 +from vllm.pooling_params import PoolingTask

29 from vllm.sequence import IntermediateTensors

31 from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only

@@ -91,8 +92,11 @@ class BertPooler(Pooler):

92 self.dense = nn.Linear(config.hidden_size, config.hidden_size)

93 self.activation = nn.Tanh()

94 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

95 - return self.pooling.get_pooling_params(task)

95 + def get_pooling_updates(

96 + self,

97 + task: PoolingTask,

98 + ) -> Optional[PoolingParamsUpdate]:

99 + return self.pooling.get_pooling_updates(task)

100

101 def forward(

102 self,

vllm/model_executor/models/gritlm.py

+116 -69

@@ -1,18 +1,24 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from array import array

4 +from typing import Optional, Union

6 +import numpy as np

7 import torch

8 import torch.nn as nn

9 +from typing_extensions import assert_never

11 from vllm.config import ModelConfig, VllmConfig

12 from vllm.logger import init_logger

11 -from vllm.model_executor.layers.pooler import PoolerHead, PoolerNormalize

13 +from vllm.model_executor.layers.pooler import (Pooler, PoolerHead,

14 + PoolerNormalize,

15 + PoolingParamsUpdate,

16 + build_output, get_prompt_lens,

17 + get_prompt_token_ids)

18 from vllm.model_executor.models.llama import LlamaForCausalLM

13 -from vllm.model_executor.pooling_metadata import (PoolingMetadata,

14 - PoolingTensors)

15 -from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput

19 +from vllm.model_executor.pooling_metadata import PoolingMetadata

20 +from vllm.pooling_params import PoolingTask

21 +from vllm.sequence import PoolerOutput

22 from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config

24 from .interfaces import SupportsV0Only

@@ -20,7 +26,8 @@ from .interfaces import SupportsV0Only

26 logger = init_logger(__name__)

23 -class GritLMPooler(nn.Module):

29 +class GritLMMeanPool(nn.Module):

30 + """As `MeanPool`, but only includes non-instruction tokens."""

32 def __init__(self, model_config: ModelConfig):

33 super().__init__()

@@ -38,8 +45,8 @@ class GritLMPooler(nn.Module):

45 for tok in ["<s>", "▁<", "<", "|", "embed", ">", "<0x0A>", "user"]

46 }

41 - def tokens_to_ids(tokens: list[str]) -> array:

42 - return array("i", [self.token_ids[token] for token in tokens])

48 + def tokens_to_ids(tokens: list[str]) -> np.ndarray:

49 + return np.array([self.token_ids[token] for token in tokens])

51 self.user_pattern_ids = tokens_to_ids(

52 ["▁<", "|", "user", "|", ">", "<0x0A>"])

@@ -48,32 +55,44 @@ class GritLMPooler(nn.Module):

55 self.embed_pattern_ids = tokens_to_ids(

56 ["▁<", "|", "embed", "|", ">", "<0x0A>"])

51 - self.head = PoolerHead(PoolerNormalize())

52 -

53 - def _find_array(self, arr: array, target: array, start_idx: int) -> int:

58 + def _find_array(

59 + self,

60 + arr: np.ndarray,

61 + target: np.ndarray,

62 + start_idx: int = 0,

63 + end_idx: Optional[int] = None,

64 + ) -> int:

65 """

55 - Find the first occurrence of target in arr starting from start_idx.

66 + Find the first occurrence of `target` in `arr` starting from

67 + `start_idx`.

69 Args:

58 - arr: The array to search within

59 - target: The consecutive subsequence to find

60 - start_idx: The starting index to search from

70 + arr: The array to search within.

71 + target: The consecutive subsequence to find.

72 + start_idx: The starting index to search from (inclusive).

73 + end_idx: The ending index to search from (exclusive).

75 Returns:

63 - int: The index of the first occurrence of target in arr.

76 + The index of the first occurrence of `target` in `arr`.

77 """

78 if start_idx < 0:

66 - raise ValueError("start_idx must be non-negative")

67 - if not target or not arr:

68 - raise ValueError("Empty arr or target not allowed")

79 + raise ValueError("`start_idx` must be non-negative")

80 + if len(arr) == 0 or len(target) == 0:

81 + raise ValueError("Empty `arr` or `target` not allowed")

83 + arr_len = len(arr)

84 target_len = len(target)

71 - for i in range(start_idx, len(arr) - target_len + 1):

72 - if arr[i:i + target_len] == target:

85 +

86 + if end_idx is None:

87 + end_idx = arr_len

88 +

89 + for i in range(start_idx, min(end_idx, arr_len - target_len + 1)):

90 + if (arr[i:i + target_len] == target).all():

91 return i

92 +

93 return -1

76 - def _get_instruction_len(self, prompt_token_ids: array) -> int:

95 + def _get_instruction_len(self, prompt_token_ids: np.ndarray) -> int:

96 """

97 Get the length of the instruction in the prompt.

@@ -83,7 +102,6 @@ class GritLMPooler(nn.Module):

102 The pattern matching is done using integers instead of strings

103 because the prompt is given as a list of token IDs.

104 """

86 -

105 instruction_len = 0

106

107 # Return no instruction in case of missing BOS token.

@@ -98,7 +116,8 @@ class GritLMPooler(nn.Module):

116 embed_pattern_ids = self.embed_pattern_ids

117 if self._find_array(prompt_token_ids,

118 self.user_pattern_ids,

101 - start_idx=1) == 1:

119 + start_idx=1,

120 + end_idx=2) == 1:

121 embed_pattern_ids = self.embed_newline_pattern_ids

122

123 # Find the embed pattern in the prompt.

@@ -116,64 +135,92 @@ class GritLMPooler(nn.Module):

135

136 return instruction_len

137

119 - def forward(

138 + def get_pooling_updates(

139 + self,

140 + task: PoolingTask,

141 + ) -> Optional[PoolingParamsUpdate]:

142 + # The equalities are split up to keep mypy happy

143 + if task == "encode" or task == "embed":

144 + return PoolingParamsUpdate(requires_token_ids=True)

145 +

146 + if task == "classify" or task == "score":

147 + return None

148 +

149 + assert_never(task)

150 +

151 + def forward_one(

152 self,

153 hidden_states: torch.Tensor,

122 - pooling_metadata: PoolingMetadata,

123 - ) -> PoolerOutput:

124 - """

125 - Pool the hidden states by summing the embeddings of

126 - non-instruction tokens.

127 - """

128 - prompts_token_ids = [

129 - token_ids.prompt_token_ids_array

130 - for _, token_ids in pooling_metadata.seq_data.items()

131 - ]

154 + prompt_len: Optional[torch.Tensor] = None,

155 + instr_len: Optional[torch.Tensor] = None,

156 + ) -> torch.Tensor:

157 + assert prompt_len is None or prompt_len == hidden_states.shape[0], \

158 + "partial prefill not supported with MEAN pooling"

159 +

160 + return hidden_states[instr_len:].mean(dim=0, dtype=torch.float32)

161 +

162 + def forward_all(

163 + self,

164 + hidden_states: torch.Tensor,

165 + prompt_lens: torch.Tensor,

166 + instr_lens: torch.Tensor,

167 + ) -> Union[list[torch.Tensor], torch.Tensor]:

168 + offset = 0

169 + pooled_data = list[torch.Tensor]()

170 +

171 + for prompt_len, instr_len in zip(prompt_lens, instr_lens):

172 + pooled_data.append(hidden_states[offset + instr_len:offset +

173 + prompt_len].mean(

174 + dim=0, dtype=torch.float32))

175 + offset += prompt_len

176

133 - instruction_lens = torch.tensor(

177 + return pooled_data

178 +

179 + def forward(

180 + self,

181 + hidden_states: Union[torch.Tensor, list[torch.Tensor]],

182 + pooling_metadata: PoolingMetadata,

183 + ) -> Union[list[torch.Tensor], torch.Tensor]:

184 + prompt_lens = get_prompt_lens(hidden_states, pooling_metadata)

185 + instr_lens = torch.tensor(

186 [

135 - self._get_instruction_len(prompt_token_ids)

136 - for prompt_token_ids in prompts_token_ids

187 + self._get_instruction_len(token_ids.cpu().numpy())

188 + for token_ids in get_prompt_token_ids(pooling_metadata)

189 ],

138 - device=hidden_states.device,

190 + device=prompt_lens.device,

191 )

192

141 - prompt_lens = PoolingTensors.from_pooling_metadata(

142 - pooling_metadata, hidden_states.device).prompt_lens

143 -

144 - mask = torch.zeros_like(hidden_states, dtype=torch.bool)

145 -

146 - start_idx = 0

147 - for prompt_len, instruction_len in zip(prompt_lens, instruction_lens):

148 - end_idx = start_idx + prompt_len

149 - mask[start_idx + instruction_len:end_idx] = True

150 - start_idx = end_idx

193 + if isinstance(hidden_states, list):

194 + return [

195 + self.forward_one(h, prompt_len, instr_len) for h, prompt_len,

196 + instr_len in zip(hidden_states, prompt_lens, instr_lens)

197 + ]

198

152 - masked_hidden_states = hidden_states.masked_fill(~mask, 0.0)

199 + return self.forward_all(hidden_states, prompt_lens, instr_lens)

200

154 - sum_embeddings = torch.zeros(len(prompt_lens),

155 - hidden_states.size(1),

156 - device=hidden_states.device)

201

158 - start_idx = 0

159 - for i, prompt_len in enumerate(prompt_lens):

160 - end_idx = start_idx + prompt_len

161 - sum_embeddings[i] = masked_hidden_states[start_idx:end_idx].sum(

162 - dim=0)

163 - start_idx = end_idx

202 +class GritLMPooler(Pooler):

203

165 - num_non_instruction_tokens = prompt_lens - instruction_lens

166 - mean_embeddings = sum_embeddings / num_non_instruction_tokens.unsqueeze(

167 - 1)

204 + def __init__(self, model_config: ModelConfig):

205 + super().__init__()

206

169 - pooled_data = self.head(mean_embeddings,

170 - pooling_metadata=pooling_metadata)

207 + self.pooling = GritLMMeanPool(model_config)

208 + self.head = PoolerHead(PoolerNormalize())

209

172 - pooled_outputs = [

173 - PoolingSequenceGroupOutput(data) for data in pooled_data

174 - ]

210 + def get_pooling_updates(

211 + self,

212 + task: PoolingTask,

213 + ) -> Optional[PoolingParamsUpdate]:

214 + return self.pooling.get_pooling_updates(task)

215

176 - return PoolerOutput(outputs=pooled_outputs)

216 + def forward(

217 + self,

218 + hidden_states: torch.Tensor,

219 + pooling_metadata: PoolingMetadata,

220 + ) -> PoolerOutput:

221 + pooled_data = self.pooling(hidden_states, pooling_metadata)

222 + pooled_data = self.head(pooled_data, pooling_metadata)

223 + return build_output(pooled_data)

224

225

226 class GritLM(LlamaForCausalLM, SupportsV0Only):

@@ -202,7 +249,7 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):

249 prefix: str = "",

250 **kwargs,

251 ) -> None:

205 - # Use full attention for pooling

252 + # Use full attention for pooling (this is why V1 is not supported yet)

253 if vllm_config.model_config.runner_type == "pooling":

254 hf_config = vllm_config.model_config.hf_config

255 hf_config.is_causal = False

vllm/model_executor/models/interfaces.py

-7

@@ -599,13 +599,6 @@ def supports_cross_encoding(

599 return is_pooling_model(model) and _supports_cross_encoding(model)

600

601

602 -def has_step_pooler(model: Union[type[object], object]) -> bool:

603 - """Check if the model uses step pooler."""

604 - from vllm.model_executor.layers.pooler import StepPooler

605 -

606 - return is_pooling_model(model) and isinstance(model.pooler, StepPooler)

607 -

608 -

602 class SupportsQuant:

603 """The interface required for all models that support quantization."""

604

vllm/model_executor/models/modernbert.py

+8 -4

@@ -14,14 +14,15 @@ from vllm.distributed import get_tensor_model_parallel_world_size

14 from vllm.model_executor.layers.linear import (QKVParallelLinear,

15 RowParallelLinear)

16 from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,

17 - PoolingMethod, PoolingTask,

17 + PoolingMethod,

18 + PoolingParamsUpdate,

19 PoolingType)

20 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

21 from vllm.model_executor.layers.vocab_parallel_embedding import (

22 VocabParallelEmbedding)

23 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

24 from vllm.model_executor.pooling_metadata import PoolingMetadata

24 -from vllm.pooling_params import PoolingParams

25 +from vllm.pooling_params import PoolingTask

26 from vllm.sequence import IntermediateTensors

28 from .interfaces import SupportsCrossEncoding, SupportsV0Only

@@ -270,8 +271,11 @@ class ModernBertPooler(Pooler):

271 eps=config.norm_eps,

272 bias=config.norm_bias)

273

273 - def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:

274 - return self.pooling.get_pooling_params(task)

274 + def get_pooling_updates(

275 + self,

276 + task: PoolingTask,

277 + ) -> Optional[PoolingParamsUpdate]:

278 + return self.pooling.get_pooling_updates(task)

279

280 def forward(

281 self,

vllm/pooling_params.py

+25 -16

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -from typing import TYPE_CHECKING, Optional

4 +from typing import TYPE_CHECKING, Literal, Optional

6 import msgspec

@@ -10,12 +10,14 @@ from vllm.sampling_params import RequestOutputKind

10 if TYPE_CHECKING:

11 from vllm.config import ModelConfig

13 +PoolingTask = Literal["encode", "embed", "classify", "score"]

14 +

16 class PoolingParams(

17 msgspec.Struct,

18 omit_defaults=True, # type: ignore[call-arg]

19 array_like=True): # type: ignore[call-arg]

18 - """API parameters for pooling models. This

20 + """API parameters for pooling models.

22 Attributes:

23 dimensions: Reduce the dimensions of embeddings

@@ -24,24 +26,33 @@ class PoolingParams(

27 dimensions: Optional[int] = None

27 - use_cross_encoder: bool = False

28 - """Internal use only."""

29 + output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY

30 - logits_processing_needs_token_ids: bool = False

31 + task: Optional[PoolingTask] = None

32 """Internal use only."""

33 - output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY

34 + requires_token_ids: bool = False

35 + """Internal use only."""

37 def clone(self) -> "PoolingParams":

38 """Returns a deep copy of the PoolingParams instance."""

39 return PoolingParams(

40 dimensions=self.dimensions,

39 - use_cross_encoder=self.use_cross_encoder,

40 - logits_processing_needs_token_ids=self.

41 - logits_processing_needs_token_ids,

41 + task=self.task,

42 + requires_token_ids=self.requires_token_ids,

43 )

44 - def verify(self, model_config: "ModelConfig") -> None:

45 + def verify(self, task: PoolingTask, model_config: "ModelConfig") -> None:

46 + if self.task is None:

47 + self.task = task

48 + elif self.task != task:

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

50 + raise ValueError(msg)

51 +

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

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

54 + # in this method

55 +

56 if self.dimensions is not None:

57 if not model_config.is_matryoshka:

58 raise ValueError(

@@ -61,12 +72,10 @@ class PoolingParams(

72 raise ValueError("Dimensions must be greater than 0")

74 def __repr__(self) -> str:

64 - return (

65 - f"PoolingParams("

66 - f"dimensions={self.dimensions}, "

67 - f"use_cross_encoder={self.use_cross_encoder}, "

68 - f"logits_processing_needs_token_ids={self.logits_processing_needs_token_ids})"

69 - )

75 + return (f"PoolingParams("

76 + f"dimensions={self.dimensions}, "

77 + f"task={self.task}, "

78 + f"requires_token_ids={self.requires_token_ids})")

80 def __post_init__(self) -> None:

81 assert self.output_kind == RequestOutputKind.FINAL_ONLY,\

vllm/v1/engine/core.py

@@ -181,6 +181,12 @@ class EngineCore:

181

182 def add_request(self, request: EngineCoreRequest):

183 """Add request to the scheduler."""

184 + if pooling_params := request.pooling_params:

185 + supported_pooling_tasks = (

186 + self.model_executor.supported_pooling_tasks)

187 + if pooling_params.task not in supported_pooling_tasks:

188 + raise ValueError(f"Unsupported task: {pooling_params.task!r} "

189 + f"Supported tasks: {supported_pooling_tasks}")

190

191 if request.mm_hashes is not None:

192 # Here, if hash exists for a multimodal input, then it will be

vllm/v1/worker/cpu_model_runner.py

-4

@@ -8,7 +8,6 @@ import torch

8 from vllm.config import VllmConfig

9 from vllm.logger import init_logger

10 from vllm.model_executor.model_loader import get_model

11 -from vllm.model_executor.models.interfaces import has_step_pooler

11 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

13 logger = init_logger(__name__)

@@ -54,9 +53,6 @@ class CPUModelRunner(GPUModelRunner):

53 logger.info("Starting to load model %s...", self.model_config.model)

54 self.model = get_model(vllm_config=self.vllm_config)

57 - if has_step_pooler(self.model):

58 - self.input_batch.logits_processing_needs_token_ids = True

59 -

56 if self.lora_config:

57 self.model = self.load_lora_model(self.model, self.model_config,

58 self.scheduler_config,

vllm/v1/worker/gpu_input_batch.py

+11 -8

@@ -70,7 +70,6 @@ class InputBatch:

70 vocab_size: int,

71 block_sizes: list[int], # The block_size of each kv cache group

72 is_spec_decode: bool = False,

73 - logits_processing_needs_token_ids: bool = False,

73 ):

74 self.is_spec_decode = is_spec_decode

75 self.max_num_reqs = max_num_reqs

@@ -79,8 +78,6 @@ class InputBatch:

78 self.device = device

79 self.pin_memory = pin_memory

80 self.vocab_size = vocab_size

82 - self.logits_processing_needs_token_ids = (

83 - logits_processing_needs_token_ids)

82 self._req_ids: list[Optional[str]] = []

83 self.req_id_to_index: dict[str, int] = {}

@@ -233,6 +230,9 @@ class InputBatch:

230 # req_index -> bad_words_token_ids

231 self.bad_words_token_ids: dict[int, list[list[int]]] = {}

232

233 + self.logits_processing_needs_token_ids = np.zeros(max_num_reqs,

234 + dtype=bool)

235 +

236 self.req_output_token_ids: list[Optional[list[int]]] = []

237

238 # This is updated each time the batch constituents change.

@@ -365,9 +365,12 @@ class InputBatch:

365 if sampling_params.bad_words_token_ids:

366 self.bad_words_token_ids[

367 req_index] = sampling_params.bad_words_token_ids

368 + elif pooling_params := request.pooling_params:

369 + self.pooling_params[req_id] = pooling_params

370 + self.logits_processing_needs_token_ids[req_index] = (

371 + pooling_params.requires_token_ids)

372 else:

369 - assert request.pooling_params is not None

370 - self.pooling_params[req_id] = request.pooling_params

373 + raise NotImplementedError(request)

374

375 # Add request lora ID

376 if request.lora_request:

@@ -620,9 +623,9 @@ class InputBatch:

623 copy_slice(self.repetition_penalties_cpu_tensor,

624 self.repetition_penalties, num_reqs)

625

623 - needs_prompt_token_ids = (not self.no_penalties or

624 - (self.num_reqs > 0

625 - and self.logits_processing_needs_token_ids))

626 + needs_prompt_token_ids = (

627 + not self.no_penalties

628 + or self.logits_processing_needs_token_ids[:num_reqs].any())

629 if needs_prompt_token_ids:

630 # The prompt tokens are used only for applying penalties or

631 # step pooling during the sampling/pooling process.

vllm/v1/worker/gpu_model_runner.py

+39 -9

@@ -4,7 +4,7 @@

4 import gc

5 import time

6 from contextlib import contextmanager

7 -from typing import TYPE_CHECKING, Any, Optional, Union

7 +from typing import TYPE_CHECKING, Any, Optional, Union, cast, get_args

9 import numpy as np

10 import torch

@@ -32,12 +32,13 @@ from vllm.logger import init_logger

32 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase

33 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

34 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

35 -from vllm.model_executor.models.interfaces import (has_step_pooler,

36 - is_mixture_of_experts)

35 +from vllm.model_executor.models.interfaces import is_mixture_of_experts

36 +from vllm.model_executor.models.interfaces_base import (VllmModelForPooling,

37 + is_pooling_model)

38 from vllm.multimodal import MULTIMODAL_REGISTRY

39 from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange

40 from vllm.multimodal.utils import group_mm_inputs_by_modality

40 -from vllm.pooling_params import PoolingParams

41 +from vllm.pooling_params import PoolingParams, PoolingTask

42 from vllm.sampling_params import SamplingType

43 from vllm.sequence import IntermediateTensors

44 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

@@ -404,6 +405,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

405 req_id = new_req_data.req_id

406 sampling_params = new_req_data.sampling_params

407 pooling_params = new_req_data.pooling_params

408 +

409 if sampling_params and \

410 sampling_params.sampling_type == SamplingType.RANDOM_SEED:

411 generator = torch.Generator(device=self.device)

@@ -411,6 +413,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):

413 else:

414 generator = None

415

416 + if pooling_params:

417 + assert pooling_params.task is not None, (

418 + "You did not set `task` in the API")

419 +

420 + model = cast(VllmModelForPooling, self.model)

421 + to_update = (model.pooler.get_pooling_updates(

422 + pooling_params.task))

423 + assert to_update is not None, (

424 + f"{pooling_params.task=} is not supported by the model")

425 +

426 + to_update.apply(pooling_params)

427 +

428 self.requests[req_id] = CachedRequestState(

429 req_id=req_id,

430 prompt_token_ids=new_req_data.prompt_token_ids,

@@ -1092,6 +1106,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1106 def get_model(self) -> nn.Module:

1107 return self.model

1108

1109 + def get_supported_pooling_tasks(self) -> list[PoolingTask]:

1110 + model = self.get_model()

1111 + if not is_pooling_model(model):

1112 + return []

1113 +

1114 + return [

1115 + task for task in get_args(PoolingTask)

1116 + if model.pooler.get_pooling_updates(task)

1117 + ]

1118 +

1119 def apply_grammar_bitmask(

1120 self,

1121 scheduler_output: "SchedulerOutput",

@@ -1737,8 +1761,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1761 )

1762 model_loader.load_weights(self.model,

1763 model_config=self.model_config)

1740 - if has_step_pooler(self.model):

1741 - self.input_batch.logits_processing_needs_token_ids = True

1764 if self.lora_config:

1765 self.model = self.load_lora_model(self.model,

1766 self.model_config,

@@ -2138,17 +2160,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2160

2161 req_num_tokens = num_tokens // num_reqs

2162

2163 + model = cast(VllmModelForPooling, self.model)

2164 + dummy_task = self.get_supported_pooling_tasks()[0]

2165 + dummy_pooling_params = PoolingParams(task=dummy_task)

2166 +

2167 + to_update = model.pooler.get_pooling_updates(dummy_task)

2168 + assert to_update is not None

2169 + to_update.apply(dummy_pooling_params)

2170 +

2171 dummy_metadata = PoolingMetadata(

2172 prompt_lens=torch.tensor([h.shape[0] for h in hidden_states_list],

2173 device=self.device),

2174 prompt_token_ids=torch.zeros((num_reqs, req_num_tokens),

2175 dtype=torch.int32,

2176 device=self.device),

2147 - pooling_params=[PoolingParams()] * num_reqs)

2177 + pooling_params=[dummy_pooling_params] * num_reqs)

2178

2179 try:

2150 - pooler_output = self.model.pooler(hidden_states=hidden_states_list,

2151 - pooling_metadata=dummy_metadata)

2180 + pooler_output = model.pooler(hidden_states=hidden_states_list,

2181 + pooling_metadata=dummy_metadata)

2182 except RuntimeError as e:

2183 if 'out of memory' in str(e):

2184 raise RuntimeError(

vllm/v1/worker/gpu_worker.py

@@ -23,6 +23,7 @@ from vllm.logger import init_logger

23 from vllm.lora.request import LoRARequest

24 from vllm.model_executor import set_random_seed

25 from vllm.platforms import current_platform

26 +from vllm.pooling_params import PoolingTask

27 from vllm.sequence import IntermediateTensors

28 from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling

29 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

@@ -309,6 +310,9 @@ class Worker(WorkerBase):

310 def get_model(self) -> nn.Module:

311 return self.model_runner.get_model()

312

313 + def get_supported_pooling_tasks(self) -> list[PoolingTask]:

314 + return self.model_runner.get_supported_pooling_tasks()

315 +

316 @torch.inference_mode()

317 def execute_model(

318 self,

vllm/v1/worker/tpu_model_runner.py

+13 -1

@@ -3,7 +3,7 @@

3 import bisect

4 import gc

5 import time

6 -from typing import TYPE_CHECKING, Any, Optional, cast

6 +from typing import TYPE_CHECKING, Any, Optional, cast, get_args

7 from unittest.mock import patch

9 import numpy as np

@@ -25,10 +25,12 @@ from vllm.logger import init_logger

25 from vllm.lora.layers import BaseLayerWithLoRA

26 from vllm.model_executor.model_loader import get_model_loader

27 from vllm.model_executor.model_loader.tpu import TPUModelLoader

28 +from vllm.model_executor.models.interfaces_base import is_pooling_model

29 from vllm.multimodal import MULTIMODAL_REGISTRY

30 from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,

31 PlaceholderRange)

32 from vllm.multimodal.utils import group_mm_inputs_by_modality

33 +from vllm.pooling_params import PoolingTask

34 from vllm.sequence import IntermediateTensors

35 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType, cdiv,

36 is_pin_memory_available, prev_power_of_2)

@@ -483,6 +485,16 @@ class TPUModelRunner(LoRAModelRunnerMixin):

485 def get_model(self) -> nn.Module:

486 return self.model

487

488 + def get_supported_pooling_tasks(self) -> list[PoolingTask]:

489 + model = self.get_model()

490 + if not is_pooling_model(model):

491 + return []

492 +

493 + return [

494 + task for task in get_args(PoolingTask)

495 + if model.pooler.get_pooling_updates(task)

496 + ]

497 +

498 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

499 """

500 Generates the KVCacheSpec by parsing the kv cache format from each

vllm/v1/worker/tpu_worker.py

@@ -19,6 +19,7 @@ from vllm.logger import init_logger

19 from vllm.lora.request import LoRARequest

20 from vllm.model_executor import set_random_seed

21 from vllm.platforms import current_platform

22 +from vllm.pooling_params import PoolingTask

23 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv

24 from vllm.v1.attention.backends.pallas import TPU_HEAD_SIZE_ALIGNMENT

25 from vllm.v1.core.sched.output import SchedulerOutput

@@ -275,6 +276,9 @@ class TPUWorker:

276 def get_model(self) -> nn.Module:

277 return self.model_runner.get_model()

278

279 + def get_supported_pooling_tasks(self) -> list[PoolingTask]:

280 + return self.model_runner.get_supported_pooling_tasks()

281 +

282 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

283 return self.model_runner.get_kv_cache_spec()

284

vllm/worker/model_runner_base.py

+13 -1

@@ -4,7 +4,7 @@

4 import dataclasses

5 from abc import ABC, abstractmethod

6 from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Type,

7 - TypeVar)

7 + TypeVar, get_args)

9 import torch

10 import torch.nn as nn

@@ -12,6 +12,8 @@ import torch.nn as nn

12 from vllm.config import VllmConfig

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.sampler import SamplerOutput

15 +from vllm.model_executor.models.interfaces_base import is_pooling_model

16 +from vllm.pooling_params import PoolingTask

17 from vllm.sequence import IntermediateTensors, SequenceGroupMetadata

19 if TYPE_CHECKING:

@@ -223,6 +225,16 @@ class ModelRunnerBase(ABC, Generic[T]):

225 def get_model(self) -> nn.Module:

226 raise NotImplementedError

227

228 + def get_supported_pooling_tasks(self) -> list[PoolingTask]:

229 + model = self.get_model()

230 + if not is_pooling_model(model):

231 + return []

232 +

233 + return [

234 + task for task in get_args(PoolingTask)

235 + if model.pooler.get_pooling_updates(task)

236 + ]

237 +

238 def execute_model(

239 self,

240 model_input: T,

vllm/worker/pooling_model_runner.py

+15 -1

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import dataclasses

5 -from typing import Any, Dict, List, Optional, Tuple, Type, Union

5 +from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast

7 import torch

@@ -10,6 +10,7 @@ from vllm.config import VllmConfig

10 from vllm.distributed import get_pp_group

11 from vllm.forward_context import set_forward_context

12 from vllm.logger import init_logger

13 +from vllm.model_executor.models.interfaces_base import VllmModelForPooling

14 from vllm.model_executor.pooling_metadata import PoolingMetadata

15 from vllm.multimodal import MultiModalKwargs

16 from vllm.pooling_params import PoolingParams

@@ -195,7 +196,20 @@ class PoolingModelRunner(

196 seq_groups: List[Tuple[List[int], PoolingParams]] = []

197 for i, seq_group_metadata in enumerate(seq_group_metadata_list):

198 seq_ids = list(seq_group_metadata.seq_data.keys())

199 +

200 pooling_params = seq_group_metadata.pooling_params

201 + assert pooling_params is not None

202 + assert pooling_params.task is not None, (

203 + "You did not set `task` in the API")

204 +

205 + to_update = (cast(VllmModelForPooling,

206 + self.model).pooler.get_pooling_updates(

207 + pooling_params.task))

208 + assert to_update is not None, (

209 + f"{pooling_params.task=} is not supported by the model")

210 +

211 + to_update.apply(pooling_params)

212 +

213 seq_groups.append((seq_ids, pooling_params))

214

215 seq_data: Dict[int, SequenceData] = {}