@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 90bd2ab6e3
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

[Model] Update pooling model interface (#21058)

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

Cyrus Leung committed Jul 18, 2025 at 00:05 UTC 90bd2ab6e3eb7e83d3f40d99fc23e6e43834743a
17 files changed +247 -345
tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py
+5 -10
@@ -11,11 +11,13 @@ from vllm.config import VllmConfig
11 from vllm.model_executor.layers.pooler import Pooler, PoolingType
12 from vllm.model_executor.models.gemma2 import Gemma2Model
13 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
14 -from vllm.model_executor.pooling_metadata import PoolingMetadata
15 -from vllm.sequence import IntermediateTensors, PoolerOutput
14 +from vllm.sequence import IntermediateTensors
15
16
17 class MyGemma2Embedding(nn.Module):
18 +
19 + is_pooling_model = True
20 +
21 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
22
23 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -24,7 +26,7 @@ class MyGemma2Embedding(nn.Module):
26 self.model = Gemma2Model(vllm_config=vllm_config,
27 prefix=maybe_prefix(prefix, "model"))
28
27 - self._pooler = Pooler.from_config_with_defaults(
29 + self.pooler = Pooler.from_config_with_defaults(
30 vllm_config.model_config.pooler_config,
31 pooling_type=PoolingType.LAST,
32 normalize=True,
@@ -54,13 +56,6 @@ class MyGemma2Embedding(nn.Module):
56 # Return all-zero embeddings
57 return torch.zeros_like(hidden_states)
58
57 - def pooler(
58 - self,
59 - hidden_states: torch.Tensor,
60 - pooling_metadata: PoolingMetadata,
61 - ) -> Optional[PoolerOutput]:
62 - return self._pooler(hidden_states, pooling_metadata)
63 -
59 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
60
61 weights = self.hf_to_vllm_mapper.apply(weights)
vllm/entrypoints/openai/protocol.py
+5 -29
@@ -1237,10 +1237,6 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1237 user: Optional[str] = None
1238 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1239
1240 - # --8<-- [start:embedding-pooling-params]
1241 - additional_data: Optional[Any] = None
1242 - # --8<-- [end:embedding-pooling-params]
1243 -
1240 # --8<-- [start:embedding-extra-params]
1241 add_special_tokens: bool = Field(
1242 default=True,
@@ -1259,8 +1255,7 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1255 # --8<-- [end:embedding-extra-params]
1256
1257 def to_pooling_params(self):
1262 - return PoolingParams(dimensions=self.dimensions,
1263 - additional_data=self.additional_data)
1258 + return PoolingParams(dimensions=self.dimensions)
1259
1260
1261 class EmbeddingChatRequest(OpenAIBaseModel):
@@ -1272,10 +1267,6 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1267 user: Optional[str] = None
1268 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1269
1275 - # --8<-- [start:chat-embedding-pooling-params]
1276 - additional_data: Optional[Any] = None
1277 - # --8<-- [end:chat-embedding-pooling-params]
1278 -
1270 # --8<-- [start:chat-embedding-extra-params]
1271 add_special_tokens: bool = Field(
1272 default=False,
@@ -1323,8 +1314,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1314 return data
1315
1316 def to_pooling_params(self):
1326 - return PoolingParams(dimensions=self.dimensions,
1327 - additional_data=self.additional_data)
1317 + return PoolingParams(dimensions=self.dimensions)
1318
1319
1320 EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
@@ -1340,10 +1330,6 @@ class ScoreRequest(OpenAIBaseModel):
1330 text_2: Union[list[str], str, ScoreMultiModalParam]
1331 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1332
1343 - # --8<-- [start:score-pooling-params]
1344 - additional_data: Optional[Any] = None
1345 - # --8<-- [end:score-pooling-params]
1346 -
1333 # --8<-- [start:score-extra-params]
1334
1335 mm_processor_kwargs: Optional[dict[str, Any]] = Field(
@@ -1362,8 +1348,7 @@ class ScoreRequest(OpenAIBaseModel):
1348 # --8<-- [end:score-extra-params]
1349
1350 def to_pooling_params(self, *, use_cross_encoder: bool = False):
1365 - return PoolingParams(use_cross_encoder=use_cross_encoder,
1366 - additional_data=self.additional_data)
1351 + return PoolingParams(use_cross_encoder=use_cross_encoder)
1352
1353
1354 class RerankRequest(OpenAIBaseModel):
@@ -1373,10 +1358,6 @@ class RerankRequest(OpenAIBaseModel):
1358 top_n: int = Field(default_factory=lambda: 0)
1359 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None
1360
1376 - # --8<-- [start:rerank-pooling-params]
1377 - additional_data: Optional[Any] = None
1378 - # --8<-- [end:rerank-pooling-params]
1379 -
1361 # --8<-- [start:rerank-extra-params]
1362
1363 mm_processor_kwargs: Optional[dict[str, Any]] = Field(
@@ -1395,8 +1376,7 @@ class RerankRequest(OpenAIBaseModel):
1376 # --8<-- [end:rerank-extra-params]
1377
1378 def to_pooling_params(self, *, use_cross_encoder: bool = False):
1398 - return PoolingParams(use_cross_encoder=use_cross_encoder,
1399 - additional_data=self.additional_data)
1379 + return PoolingParams(use_cross_encoder=use_cross_encoder)
1380
1381
1382 class RerankDocument(BaseModel):
@@ -1534,10 +1514,6 @@ class ClassificationRequest(OpenAIBaseModel):
1514 truncate_prompt_tokens: Optional[int] = None
1515 user: Optional[str] = None
1516
1537 - # --8<-- [start:classification-pooling-params]
1538 - additional_data: Optional[Any] = None
1539 - # --8<-- [end:classification-pooling-params]
1540 -
1517 # --8<-- [start:classification-extra-params]
1518 priority: int = Field(
1519 default=0,
@@ -1550,7 +1526,7 @@ class ClassificationRequest(OpenAIBaseModel):
1526 # --8<-- [end:classification-extra-params]
1527
1528 def to_pooling_params(self):
1553 - return PoolingParams(additional_data=self.additional_data)
1529 + return PoolingParams()
1530
1531
1532 class ClassificationData(OpenAIBaseModel):
vllm/model_executor/layers/pooler.py
+112 -64
@@ -3,22 +3,25 @@
3 from abc import ABC, abstractmethod
4 from dataclasses import dataclass
5 from enum import IntEnum
6 -from typing import Callable, Optional, TypeVar, Union
6 +from typing import Callable, Literal, Optional, TypeVar, Union
7
8 import torch
9 import torch.nn as nn
10 import torch.nn.functional as F
11 from transformers import PretrainedConfig
12 +from typing_extensions import assert_never
13
14 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
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
22
23 PoolingMetadata = Union[V0PoolingMetadata, V1PoolingMetadata]
24 +PoolingTask = Literal["encode", "embed", "classify", "score"]
25
26
27 class PoolingType(IntEnum):
@@ -64,6 +67,48 @@ class ResolvedPoolingConfig:
67 )
68
69
70 +class Pooler(nn.Module, ABC):
71 + """The interface required for all poolers used in pooling models in vLLM."""
72 +
73 + @staticmethod
74 + def from_config_with_defaults(
75 + pooler_config: PoolerConfig,
76 + pooling_type: PoolingType,
77 + normalize: bool,
78 + softmax: bool,
79 + step_tag_id: Optional[int] = None,
80 + returned_token_ids: Optional[list[int]] = None,
81 + ) -> "Pooler":
82 + resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
83 + pooler_config=pooler_config,
84 + pooling_type=pooling_type,
85 + normalize=normalize,
86 + softmax=softmax,
87 + step_tag_id=step_tag_id,
88 + returned_token_ids=returned_token_ids,
89 + )
90 +
91 + if pooling_type == PoolingType.STEP:
92 + return StepPooler.from_config(resolved_config)
93 +
94 + return SimplePooler.from_config(resolved_config)
95 +
96 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
97 + """
98 + Construct the pooling parameters to use for a task,
99 + or `None` if the task is not supported.
100 + """
101 + return None
102 +
103 + @abstractmethod
104 + def forward(
105 + self,
106 + hidden_states: Union[list[torch.Tensor], torch.Tensor],
107 + pooling_metadata: PoolingMetadata,
108 + ) -> PoolerOutput:
109 + raise NotImplementedError
110 +
111 +
112 def get_prompt_lens(
113 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
114 pooling_metadata: PoolingMetadata,
@@ -104,17 +149,6 @@ def build_output(all_data: torch.Tensor) -> PoolerOutput:
149 return PoolerOutput(outputs=all_outputs)
150
151
107 -class BasePooler(nn.Module):
108 -
109 - @abstractmethod
110 - def forward(
111 - self,
112 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
113 - pooling_metadata: PoolingMetadata,
114 - ) -> PoolerOutput:
115 - raise NotImplementedError
116 -
117 -
152 class PoolingMethod(nn.Module, ABC):
153
154 @staticmethod
@@ -130,6 +164,10 @@ class PoolingMethod(nn.Module, ABC):
164
165 raise NotImplementedError(f"Unsupported method: {pooling_type}")
166
167 + @abstractmethod
168 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
169 + raise NotImplementedError
170 +
171 @abstractmethod
172 def forward_one(
173 self,
@@ -168,6 +206,14 @@ class PoolingMethod(nn.Module, ABC):
206
207 class CLSPool(PoolingMethod):
208
209 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
210 + # The equalities are split up to keep mypy happy
211 + if (task == "encode" or task == "embed" or task == "classify"
212 + or task == "score"):
213 + return PoolingParams()
214 +
215 + assert_never(task)
216 +
217 def forward_one(
218 self,
219 hidden_states: torch.Tensor,
@@ -190,6 +236,14 @@ class CLSPool(PoolingMethod):
236
237 class LastPool(PoolingMethod):
238
239 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
240 + # The equalities are split up to keep mypy happy
241 + if (task == "encode" or task == "embed" or task == "classify"
242 + or task == "score"):
243 + return PoolingParams()
244 +
245 + assert_never(task)
246 +
247 def forward_one(
248 self,
249 hidden_states: torch.Tensor,
@@ -208,6 +262,16 @@ class LastPool(PoolingMethod):
262
263 class AllPool(PoolingMethod):
264
265 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
266 + if task == "encode":
267 + return PoolingParams()
268 +
269 + # The equalities are split up to keep mypy happy
270 + if task == "embed" or task == "classify" or task == "score":
271 + return None
272 +
273 + assert_never(task)
274 +
275 def forward_one(
276 self,
277 hidden_states: torch.Tensor,
@@ -235,6 +299,14 @@ class AllPool(PoolingMethod):
299
300 class MeanPool(PoolingMethod):
301
302 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
303 + # The equalities are split up to keep mypy happy
304 + if (task == "encode" or task == "embed" or task == "classify"
305 + or task == "score"):
306 + return PoolingParams()
307 +
308 + assert_never(task)
309 +
310 def forward_one(
311 self,
312 hidden_states: torch.Tensor,
@@ -345,25 +417,6 @@ class LambdaPoolerActivation(PoolerActivation):
417
418 class PoolerHead(nn.Module):
419
348 - @classmethod
349 - def from_config_with_defaults(
350 - cls,
351 - pooler_config: PoolerConfig,
352 - pooling_type: PoolingType,
353 - normalize: bool,
354 - softmax: bool,
355 - ) -> "PoolerHead":
356 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
357 - pooler_config=pooler_config,
358 - pooling_type=pooling_type,
359 - normalize=normalize,
360 - softmax=softmax,
361 - step_tag_id=None,
362 - returned_token_ids=None,
363 - )
364 -
365 - return cls.from_config(resolved_config)
366 -
420 @classmethod
421 def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "PoolerHead":
422 if pooler_config.normalize and pooler_config.softmax:
@@ -424,21 +477,17 @@ class PoolerHead(nn.Module):
477 return self.activation(pooled_data)
478
479
427 -class SimplePooler(BasePooler):
480 +class SimplePooler(Pooler):
481 """A layer that pools specific information from hidden states.
482
483 This layer does the following:
484 1. Extracts specific tokens or aggregates data based on pooling method.
485 2. Normalizes output if specified.
486 3. Returns structured results as `PoolerOutput`.
434 -
435 - Attributes:
436 - pooling_type: The type of pooling to use.
437 - normalize: Whether to normalize the pooled data.
487 """
488
489 @classmethod
441 - def from_config_with_defaults(
490 + def from_config_with_defaults( # type: ignore[override]
491 cls,
492 pooler_config: PoolerConfig,
493 pooling_type: PoolingType,
@@ -471,6 +520,9 @@ class SimplePooler(BasePooler):
520 self.pooling = pooling
521 self.head = head
522
523 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
524 + return self.pooling.get_pooling_params(task)
525 +
526 def forward(
527 self,
528 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
@@ -481,7 +533,7 @@ class SimplePooler(BasePooler):
533 return build_output(pooled_data)
534
535
484 -class StepPooler(BasePooler):
536 +class StepPooler(Pooler):
537
538 @classmethod
539 def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "StepPooler":
@@ -543,6 +595,16 @@ class StepPooler(BasePooler):
595
596 return pooled_data
597
598 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
599 + if task == "encode":
600 + return PoolingParams(logits_processing_needs_token_ids=True)
601 +
602 + # The equalities are split up to keep mypy happy
603 + if task == "embed" or task == "classify" or task == "score":
604 + return None
605 +
606 + assert_never(task)
607 +
608 def forward(
609 self,
610 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
@@ -553,32 +615,6 @@ class StepPooler(BasePooler):
615 return build_output(pooled_data)
616
617
556 -class Pooler(nn.Module):
557 -
558 - @staticmethod
559 - def from_config_with_defaults(
560 - pooler_config: PoolerConfig,
561 - pooling_type: PoolingType,
562 - normalize: bool,
563 - softmax: bool,
564 - step_tag_id: Optional[int] = None,
565 - returned_token_ids: Optional[list[int]] = None,
566 - ) -> BasePooler:
567 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
568 - pooler_config=pooler_config,
569 - pooling_type=pooling_type,
570 - normalize=normalize,
571 - softmax=softmax,
572 - step_tag_id=step_tag_id,
573 - returned_token_ids=returned_token_ids,
574 - )
575 -
576 - if pooling_type == PoolingType.STEP:
577 - return StepPooler.from_config(resolved_config)
578 -
579 - return SimplePooler.from_config(resolved_config)
580 -
581 -
618 PoolingFn = Callable[
619 [Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
620 Union[torch.Tensor, list[torch.Tensor]]]
@@ -618,6 +654,18 @@ class ClassifierPooler(nn.Module):
654 return (self.cross_encoder_act_fn
655 if use_cross_encoder else self.classification_act_fn)
656
657 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
658 + if task == "encode":
659 + return PoolingParams()
660 + if task == "embed":
661 + return None
662 + if task == "classify":
663 + return PoolingParams()
664 + if task == "score":
665 + return PoolingParams(use_cross_encoder=True)
666 +
667 + assert_never(task)
668 +
669 def forward(
670 self,
671 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
vllm/model_executor/models/adapters.py
+8 -23
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Iterable
5 -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union, cast
5 +from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast
6
7 import torch
8 import torch.nn as nn
@@ -42,13 +42,14 @@ def _create_pooling_model_cls(
42 default_softmax: bool,
43 ) -> _T:
44 # Lazy import
45 - from vllm.model_executor.layers.pooler import Pooler, PoolerOutput
46 - from vllm.model_executor.pooling_metadata import PoolingMetadata
45 + from vllm.model_executor.layers.pooler import Pooler
46
47 from .utils import AutoWeightsLoader, WeightsMapper
48
49 class ModelForPooling(orig_cls, VllmModelForPooling):
50
51 + is_pooling_model = True
52 +
53 def __init__(
54 self,
55 *,
@@ -66,27 +67,20 @@ def _create_pooling_model_cls(
67 delattr(self, attr)
68
69 # If the model already defines a pooler instance, don't overwrite it
69 - if not getattr(self, "_pooler", None):
70 + if not getattr(self, "pooler", None):
71 self._init_pooler(vllm_config, prefix=prefix)
72
73 def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
74 pooler_config = vllm_config.model_config.pooler_config
75 assert pooler_config is not None
76
76 - self._pooler = Pooler.from_config_with_defaults(
77 + self.pooler = Pooler.from_config_with_defaults(
78 pooler_config,
79 pooling_type=default_pooling_type,
80 normalize=default_normalize,
81 softmax=default_softmax,
82 )
83
83 - def pooler(
84 - self,
85 - hidden_states: torch.Tensor,
86 - pooling_metadata: PoolingMetadata,
87 - ) -> PoolerOutput:
88 - return self._pooler(hidden_states, pooling_metadata)
89 -
84 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
85 # TODO: Support uninitialized params tracking
86
@@ -171,10 +165,8 @@ def as_seq_cls_model(cls: _T) -> _T:
165 # Lazy import
166 from vllm.model_executor.layers.linear import RowParallelLinear
167 from vllm.model_executor.layers.pooler import (ClassifierPooler,
174 - PoolerOutput, PoolingType,
175 - SimplePooler)
168 + PoolingType, SimplePooler)
169 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
177 - from vllm.model_executor.pooling_metadata import PoolingMetadata
170 from vllm.sequence import IntermediateTensors
171
172 from .utils import maybe_prefix
@@ -213,7 +205,7 @@ def as_seq_cls_model(cls: _T) -> _T:
205 softmax=True,
206 )
207
216 - self._pooler = ClassifierPooler(
208 + self.pooler = ClassifierPooler(
209 vllm_config.model_config,
210 pooling=pooler.pooling,
211 classifier=self._classifier,
@@ -234,13 +226,6 @@ def as_seq_cls_model(cls: _T) -> _T:
226 return super().forward(input_ids, positions, intermediate_tensors,
227 inputs_embeds)
228
237 - def pooler(
238 - self,
239 - hidden_states: Union[torch.Tensor, list[torch.Tensor]],
240 - pooling_metadata: PoolingMetadata,
241 - ) -> PoolerOutput:
242 - return self._pooler(hidden_states, pooling_metadata)
243 -
229 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
230 tokens = getattr(self.config, "classifier_from_token", None)
231 method = getattr(self.config, "method", None)
vllm/model_executor/models/bert.py
+18 -19
@@ -18,12 +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, PoolingType)
21 + PoolingMethod, PoolingTask,
22 + PoolingType)
23 from vllm.model_executor.layers.quantization import QuantizationConfig
24 from vllm.model_executor.layers.vocab_parallel_embedding import (
25 VocabParallelEmbedding)
26 from vllm.model_executor.pooling_metadata import PoolingMetadata
26 -from vllm.sequence import IntermediateTensors, PoolerOutput
27 +from vllm.pooling_params import PoolingParams
28 +from vllm.sequence import IntermediateTensors
29
30 from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only
31 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
@@ -80,7 +82,7 @@ class BertEmbedding(nn.Module):
82 return embeddings
83
84
83 -class BertPooler(nn.Module):
85 +class BertPooler(Pooler):
86
87 def __init__(self, config: BertConfig):
88 super().__init__()
@@ -89,6 +91,9 @@ class BertPooler(nn.Module):
91 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
92 self.activation = nn.Tanh()
93
94 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
95 + return self.pooling.get_pooling_params(task)
96 +
97 def forward(
98 self,
99 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
@@ -319,6 +324,9 @@ class BertOutput(nn.Module):
324
325
326 class BertModel(nn.Module, SupportsQuant):
327 +
328 + is_pooling_model = True
329 +
330 packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
331
332 def __init__(self,
@@ -403,12 +411,15 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
411 _pooler: An instance of Pooler used for pooling operations.
412 """
413
414 + is_pooling_model = True
415 +
416 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
417 super().__init__()
418 +
419 pooler_config = vllm_config.model_config.pooler_config
420 self.model = self._build_model(vllm_config=vllm_config,
421 prefix=maybe_prefix(prefix, "model"))
411 - self._pooler = self._build_pooler(pooler_config)
422 + self.pooler = self._build_pooler(pooler_config)
423
424 def forward(
425 self,
@@ -422,13 +433,6 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
433 inputs_embeds=inputs_embeds,
434 intermediate_tensors=intermediate_tensors)
435
425 - def pooler(
426 - self,
427 - hidden_states: torch.Tensor,
428 - pooling_metadata: PoolingMetadata,
429 - ) -> Optional[PoolerOutput]:
430 - return self._pooler(hidden_states, pooling_metadata)
431 -
436 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
437 weights_list = list(weights)
438
@@ -466,6 +470,8 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
470 _pooler: An instance of Pooler used for pooling operations.
471 """
472
473 + is_pooling_model = True
474 +
475 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
476 super().__init__()
477 config = vllm_config.model_config.hf_config
@@ -476,7 +482,7 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
482 embedding_class=BertEmbedding,
483 add_pooling_layer=True)
484 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
479 - self._pooler = ClassifierPooler(
485 + self.pooler = ClassifierPooler(
486 vllm_config.model_config,
487 pooling=self.bert.pooler,
488 classifier=self.classifier,
@@ -487,13 +493,6 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
493 loaded_params = loader.load_weights(weights)
494 return loaded_params
495
490 - def pooler(
491 - self,
492 - hidden_states: torch.Tensor,
493 - pooling_metadata: PoolingMetadata,
494 - ) -> Optional[PoolerOutput]:
495 - return self._pooler(hidden_states, pooling_metadata)
496 -
496 def forward(
497 self,
498 input_ids: Optional[torch.Tensor],
vllm/model_executor/models/gpt2.py
+4 -10
@@ -40,9 +40,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
40 from vllm.model_executor.layers.vocab_parallel_embedding import (
41 ParallelLMHead, VocabParallelEmbedding)
42 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43 -from vllm.model_executor.pooling_metadata import PoolingMetadata
43 from vllm.model_executor.sampling_metadata import SamplingMetadata
45 -from vllm.sequence import IntermediateTensors, PoolerOutput
44 +from vllm.sequence import IntermediateTensors
45
46 from ..layers.pooler import Pooler, PoolingType
47 from .interfaces import SupportsPP
@@ -332,6 +331,8 @@ class GPT2ForSequenceClassification(nn.Module):
331 _pooler: An instance of Pooler used for pooling operations.
332 """
333
334 + is_pooling_model = True
335 +
336 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
337 super().__init__()
338 config = vllm_config.model_config.hf_config
@@ -339,7 +340,7 @@ class GPT2ForSequenceClassification(nn.Module):
340 prefix=maybe_prefix(prefix, "gpt2"))
341 self.score = nn.Linear(config.n_embd, config.num_labels, bias=False)
342 pooler_config = vllm_config.model_config.pooler_config
342 - self._pooler = Pooler.from_config_with_defaults(
343 + self.pooler = Pooler.from_config_with_defaults(
344 pooler_config,
345 pooling_type=PoolingType.LAST,
346 normalize=False,
@@ -349,13 +350,6 @@ class GPT2ForSequenceClassification(nn.Module):
350 loader = AutoWeightsLoader(self)
351 return loader.load_weights(weights)
352
352 - def pooler(
353 - self,
354 - hidden_states: torch.Tensor,
355 - pooling_metadata: PoolingMetadata,
356 - ) -> Optional[PoolerOutput]:
357 - return self._pooler(hidden_states, pooling_metadata)
358 -
353 def forward(
354 self,
355 input_ids: torch.Tensor,
vllm/model_executor/models/gritlm.py
+3 -9
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from array import array
5 -from typing import Optional
5
6 import torch
7 import torch.nn as nn
@@ -195,6 +194,8 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
194 - "<|user|>\nPROMPT\n<|assistant|>\n"
195 """
196
197 + is_pooling_model = True
198 +
199 def __init__(
200 self,
201 vllm_config: VllmConfig,
@@ -214,11 +215,4 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
215
216 super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
217
217 - self._pooler = GritLMPooler(vllm_config.model_config)
218 -
219 - def pooler(
220 - self,
221 - hidden_states: torch.Tensor,
222 - pooling_metadata: PoolingMetadata,
223 - ) -> Optional[PoolerOutput]:
224 - return self._pooler(hidden_states, pooling_metadata)
218 + self.pooler = GritLMPooler(vllm_config.model_config)
vllm/model_executor/models/interfaces.py
+12 -74
@@ -119,13 +119,6 @@ class SupportsMultiModal(Protocol):
119 ...
120
121
122 -# We can't use runtime_checkable with ClassVar for issubclass checks
123 -# so we need to treat the class as an instance and use isinstance instead
124 -@runtime_checkable
125 -class _SupportsMultiModalType(Protocol):
126 - supports_multimodal: Literal[True]
127 -
128 -
122 @overload
123 def supports_multimodal(
124 model: type[object]) -> TypeIs[type[SupportsMultiModal]]:
@@ -140,10 +133,7 @@ def supports_multimodal(model: object) -> TypeIs[SupportsMultiModal]:
133 def supports_multimodal(
134 model: Union[type[object], object],
135 ) -> Union[TypeIs[type[SupportsMultiModal]], TypeIs[SupportsMultiModal]]:
143 - if isinstance(model, type):
144 - return isinstance(model, _SupportsMultiModalType)
145 -
146 - return isinstance(model, SupportsMultiModal)
136 + return getattr(model, "supports_multimodal", False)
137
138
139 @runtime_checkable
@@ -174,13 +164,6 @@ class SupportsScoreTemplate(Protocol):
164 ...
165
166
177 -# We can't use runtime_checkable with ClassVar for issubclass checks
178 -# so we need to treat the class as an instance and use isinstance instead
179 -@runtime_checkable
180 -class _SupportsScoreTemplateType(Protocol):
181 - supports_score_template: Literal[True]
182 -
183 -
167 @overload
168 def supports_score_template(
169 model: type[object]) -> TypeIs[type[SupportsScoreTemplate]]:
@@ -195,11 +178,7 @@ def supports_score_template(model: object) -> TypeIs[SupportsScoreTemplate]:
178 def supports_score_template(
179 model: Union[type[object], object],
180 ) -> Union[TypeIs[type[SupportsScoreTemplate]], TypeIs[SupportsScoreTemplate]]:
198 -
199 - if isinstance(model, type):
200 - return isinstance(model, _SupportsScoreTemplateType)
201 -
202 - return isinstance(model, SupportsScoreTemplate)
181 + return getattr(model, "supports_score_template", False)
182
183
184 @runtime_checkable
@@ -409,11 +388,6 @@ class HasInnerState(Protocol):
388 """
389
390
412 -@runtime_checkable
413 -class _HasInnerStateType(Protocol):
414 - has_inner_state: ClassVar[Literal[True]]
415 -
416 -
391 @overload
392 def has_inner_state(model: object) -> TypeIs[HasInnerState]:
393 ...
@@ -427,10 +401,7 @@ def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]:
401 def has_inner_state(
402 model: Union[type[object], object]
403 ) -> Union[TypeIs[type[HasInnerState]], TypeIs[HasInnerState]]:
430 - if isinstance(model, type):
431 - return isinstance(model, _HasInnerStateType)
432 -
433 - return isinstance(model, HasInnerState)
404 + return getattr(model, "has_inner_state", False)
405
406
407 @runtime_checkable
@@ -446,11 +417,6 @@ class IsAttentionFree(Protocol):
417 """
418
419
449 -@runtime_checkable
450 -class _IsAttentionFreeType(Protocol):
451 - is_attention_free: ClassVar[Literal[True]]
452 -
453 -
420 @overload
421 def is_attention_free(model: object) -> TypeIs[IsAttentionFree]:
422 ...
@@ -464,10 +430,7 @@ def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]:
430 def is_attention_free(
431 model: Union[type[object], object]
432 ) -> Union[TypeIs[type[IsAttentionFree]], TypeIs[IsAttentionFree]]:
467 - if isinstance(model, type):
468 - return isinstance(model, _IsAttentionFreeType)
469 -
470 - return isinstance(model, IsAttentionFree)
433 + return getattr(model, "is_attention_free", False)
434
435
436 @runtime_checkable
@@ -502,11 +465,6 @@ class IsHybrid(Protocol):
465 ...
466
467
505 -@runtime_checkable
506 -class _IsHybridType(Protocol):
507 - is_hybrid: ClassVar[Literal[True]]
508 -
509 -
468 @overload
469 def is_hybrid(model: object) -> TypeIs[IsHybrid]:
470 ...
@@ -520,10 +478,7 @@ def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]:
478 def is_hybrid(
479 model: Union[type[object], object]
480 ) -> Union[TypeIs[type[IsHybrid]], TypeIs[IsHybrid]]:
523 - if isinstance(model, type):
524 - return isinstance(model, _IsHybridType)
525 -
526 - return isinstance(model, IsHybrid)
481 + return getattr(model, "is_hybrid", False)
482
483
484 @runtime_checkable
@@ -598,11 +553,6 @@ class HasNoOps(Protocol):
553 has_noops: ClassVar[Literal[True]] = True
554
555
601 -@runtime_checkable
602 -class _HasNoOpsType(Protocol):
603 - has_noops: ClassVar[Literal[True]]
604 -
605 -
556 @overload
557 def has_noops(model: object) -> TypeIs[HasNoOps]:
558 ...
@@ -616,10 +566,7 @@ def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]:
566 def has_noops(
567 model: Union[type[object], object]
568 ) -> Union[TypeIs[type[HasNoOps]], TypeIs[HasNoOps]]:
619 - if isinstance(model, type):
620 - return isinstance(model, _HasNoOpsType)
621 -
622 - return isinstance(model, HasNoOps)
569 + return getattr(model, "has_noops", False)
570
571
572 @runtime_checkable
@@ -643,11 +590,7 @@ def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]:
590 def _supports_cross_encoding(
591 model: Union[type[object], object],
592 ) -> Union[TypeIs[type[SupportsCrossEncoding]], TypeIs[SupportsCrossEncoding]]:
646 -
647 - if isinstance(model, type):
648 - return isinstance(model, SupportsCrossEncoding)
649 -
650 - return isinstance(model, SupportsCrossEncoding)
593 + return getattr(model, "supports_cross_encoding", False)
594
595
596 def supports_cross_encoding(
@@ -658,8 +601,9 @@ def supports_cross_encoding(
601
602 def has_step_pooler(model: Union[type[object], object]) -> bool:
603 """Check if the model uses step pooler."""
661 - return is_pooling_model(model) and any(
662 - type(module).__name__ == "StepPooler" for module in model.modules())
604 + from vllm.model_executor.layers.pooler import StepPooler
605 +
606 + return is_pooling_model(model) and isinstance(model.pooler, StepPooler)
607
608
609 class SupportsQuant:
@@ -770,10 +714,7 @@ def supports_transcription(model: object) -> TypeIs[SupportsTranscription]:
714 def supports_transcription(
715 model: Union[type[object], object],
716 ) -> Union[TypeIs[type[SupportsTranscription]], TypeIs[SupportsTranscription]]:
773 - if isinstance(model, type):
774 - return isinstance(model, SupportsTranscription)
775 -
776 - return isinstance(model, SupportsTranscription)
717 + return getattr(model, "supports_transcription", False)
718
719
720 @runtime_checkable
@@ -796,7 +737,4 @@ def supports_v0_only(model: object) -> TypeIs[SupportsV0Only]:
737 def supports_v0_only(
738 model: Union[type[object], object],
739 ) -> Union[TypeIs[type[SupportsV0Only]], TypeIs[SupportsV0Only]]:
799 - if isinstance(model, type):
800 - return isinstance(model, SupportsV0Only)
801 -
802 - return isinstance(model, SupportsV0Only)
740 + return getattr(model, "supports_v0_only", False)
vllm/model_executor/models/interfaces_base.py
+16 -17
@@ -1,8 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import (TYPE_CHECKING, Optional, Protocol, Union, overload,
5 - runtime_checkable)
3 +from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,
4 + Union, overload, runtime_checkable)
5
6 import torch
7 import torch.nn as nn
@@ -13,8 +12,7 @@ from vllm.utils import supports_kw
12
13 if TYPE_CHECKING:
14 from vllm.config import VllmConfig
16 - from vllm.model_executor.layers.pooler import PoolerOutput
17 - from vllm.model_executor.pooling_metadata import PoolingMetadata
15 + from vllm.model_executor.layers.pooler import Pooler
16 from vllm.model_executor.sampling_metadata import SamplingMetadata
17
18 logger = init_logger(__name__)
@@ -130,16 +128,20 @@ def is_text_generation_model(
128
129
130 @runtime_checkable
133 -class VllmModelForPooling(VllmModel[T], Protocol[T]):
131 +class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
132 """The interface required for all pooling models in vLLM."""
133
136 - def pooler(
137 - self,
138 - hidden_states: T,
139 - pooling_metadata: "PoolingMetadata",
140 - ) -> "PoolerOutput":
141 - """Only called on TP rank 0."""
142 - ...
134 + is_pooling_model: ClassVar[Literal[True]] = True
135 + """
136 + A flag that indicates this model supports pooling.
137 +
138 + Note:
139 + There is no need to redefine this flag if this class is in the
140 + MRO of your model class.
141 + """
142 +
143 + pooler: "Pooler"
144 + """The pooler is only called on TP rank 0."""
145
146
147 @overload
@@ -158,7 +160,4 @@ def is_pooling_model(
160 if not is_vllm_model(model):
161 return False
162
161 - if isinstance(model, type):
162 - return isinstance(model, VllmModelForPooling)
163 -
164 - return isinstance(model, VllmModelForPooling)
163 + return getattr(model, "is_pooling_model", False)
vllm/model_executor/models/internlm2.py
+4 -10
@@ -28,9 +28,8 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
28 from vllm.model_executor.layers.vocab_parallel_embedding import (
29 ParallelLMHead, VocabParallelEmbedding)
30 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
31 -from vllm.model_executor.pooling_metadata import PoolingMetadata
31 from vllm.model_executor.sampling_metadata import SamplingMetadata
33 -from vllm.sequence import IntermediateTensors, PoolerOutput
32 +from vllm.sequence import IntermediateTensors
33
34 from .interfaces import SupportsLoRA, SupportsPP
35 from .utils import (is_pp_missing_parameter,
@@ -404,6 +403,8 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
403
404 class InternLM2ForRewardModel(InternLM2ForCausalLM):
405
406 + is_pooling_model = True
407 +
408 def __init__(
409 self,
410 *,
@@ -428,7 +429,7 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
429 )
430
431 pooler_config = vllm_config.model_config.pooler_config
431 - self._pooler = Pooler.from_config_with_defaults(
432 + self.pooler = Pooler.from_config_with_defaults(
433 pooler_config,
434 pooling_type=PoolingType.ALL,
435 normalize=False,
@@ -446,10 +447,3 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
447 inputs_embeds)
448 logits, _ = self.v_head(hidden_states)
449 return logits
449 -
450 - def pooler(
451 - self,
452 - hidden_states: torch.Tensor,
453 - pooling_metadata: PoolingMetadata,
454 - ) -> Optional[PoolerOutput]:
455 - return self._pooler(hidden_states, pooling_metadata)
vllm/model_executor/models/jamba.py
+4 -10
@@ -27,9 +27,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28 from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
29 MambaCacheParams)
30 -from vllm.model_executor.pooling_metadata import PoolingMetadata
30 from vllm.model_executor.sampling_metadata import SamplingMetadata
32 -from vllm.sequence import IntermediateTensors, PoolerOutput
31 +from vllm.sequence import IntermediateTensors
32 from vllm.utils import LayerBlockType
33
34 from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
@@ -563,6 +562,8 @@ def _is_moe_layer(name: str):
562
563 class JambaForSequenceClassification(JambaForCausalLM):
564
565 + is_pooling_model = True
566 +
567 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
568 super().__init__(vllm_config=vllm_config, prefix=prefix)
569
@@ -590,16 +591,9 @@ class JambaForSequenceClassification(JambaForCausalLM):
591 softmax=False,
592 )
593
593 - self._pooler = ClassifierPooler(
594 + self.pooler = ClassifierPooler(
595 vllm_config.model_config,
596 pooling=pooler.pooling,
597 classifier=self.score,
598 act_fn=pooler.head.activation,
599 )
599 -
600 - def pooler(
601 - self,
602 - hidden_states: torch.Tensor,
603 - pooling_metadata: PoolingMetadata,
604 - ) -> Optional[PoolerOutput]:
605 - return self._pooler(hidden_states, pooling_metadata)
vllm/model_executor/models/jina_vl.py
+4 -11
@@ -13,9 +13,8 @@ from vllm.logger import init_logger
13 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
14 RowParallelLinear)
15 from vllm.model_executor.layers.pooler import Pooler, PoolingType
16 -from vllm.model_executor.pooling_metadata import PoolingMetadata
16 from vllm.multimodal import MULTIMODAL_REGISTRY
18 -from vllm.sequence import IntermediateTensors, PoolerOutput
17 +from vllm.sequence import IntermediateTensors
18
19 from .interfaces import (SupportsCrossEncoding, SupportsMultiModal,
20 SupportsScoreTemplate)
@@ -72,6 +71,8 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
71 SupportsCrossEncoding,
72 SupportsMultiModal,
73 SupportsScoreTemplate):
74 +
75 + is_pooling_model = True
76 weight_mapper = WeightsMapper(
77 orig_to_new_prefix={
78 "score.0.": "score.dense.",
@@ -95,7 +96,7 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
96
97 self.score = JinaVLScorer(config)
98
98 - self._pooler = Pooler.from_config_with_defaults(
99 + self.pooler = Pooler.from_config_with_defaults(
100 pooler_config,
101 pooling_type=PoolingType.LAST,
102 normalize=False,
@@ -137,14 +138,6 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
138 logits = self.score(hidden_states) - self.LOGIT_BIAS
139 return logits
140
140 - def pooler(
141 - self,
142 - hidden_states: torch.Tensor,
143 - pooling_metadata: PoolingMetadata,
144 - ) -> Optional[PoolerOutput]:
145 - return self._pooler(hidden_states, pooling_metadata)
146 -
141 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
148 -
142 loader = AutoWeightsLoader(self)
143 return loader.load_weights(weights, mapper=self.weight_mapper)
vllm/model_executor/models/modernbert.py
+12 -12
@@ -13,14 +13,16 @@ from vllm.config import VllmConfig
13 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 (BasePooler, ClassifierPooler,
17 - PoolingMethod, PoolingType)
16 +from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,
17 + PoolingMethod, PoolingTask,
18 + PoolingType)
19 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
20 from vllm.model_executor.layers.vocab_parallel_embedding import (
21 VocabParallelEmbedding)
22 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23 from vllm.model_executor.pooling_metadata import PoolingMetadata
23 -from vllm.sequence import IntermediateTensors, PoolerOutput
24 +from vllm.pooling_params import PoolingParams
25 +from vllm.sequence import IntermediateTensors
26
27 from .interfaces import SupportsCrossEncoding, SupportsV0Only
28 from .utils import WeightsMapper, maybe_prefix
@@ -253,7 +255,7 @@ class ModernBertModel(nn.Module):
255 return norm_outputs
256
257
256 -class ModernBertPooler(BasePooler):
258 +class ModernBertPooler(Pooler):
259
260 def __init__(self, config: ModernBertConfig):
261 super().__init__()
@@ -268,6 +270,9 @@ class ModernBertPooler(BasePooler):
270 eps=config.norm_eps,
271 bias=config.norm_bias)
272
273 + def get_pooling_params(self, task: PoolingTask) -> Optional[PoolingParams]:
274 + return self.pooling.get_pooling_params(task)
275 +
276 def forward(
277 self,
278 hidden_states: Union[torch.Tensor, list[torch.Tensor]],
@@ -281,6 +286,8 @@ class ModernBertPooler(BasePooler):
286 class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
287 SupportsCrossEncoding):
288
289 + is_pooling_model = True
290 +
291 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
292 super().__init__()
293 config = vllm_config.model_config.hf_config
@@ -288,7 +295,7 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
295 self.model = ModernBertModel(vllm_config=vllm_config,
296 prefix=maybe_prefix(prefix, "modernbert"))
297 self.classifier = nn.Linear(config.hidden_size, config.num_labels)
291 - self._pooler = ClassifierPooler(
298 + self.pooler = ClassifierPooler(
299 vllm_config.model_config,
300 pooling=ModernBertPooler(config),
301 classifier=self.classifier,
@@ -321,13 +328,6 @@ class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
328 default_weight_loader)
329 weight_loader(param, loaded_weight)
330
324 - def pooler(
325 - self,
326 - hidden_states: torch.Tensor,
327 - pooling_metadata: PoolingMetadata,
328 - ) -> Optional[PoolerOutput]:
329 - return self._pooler(hidden_states, pooling_metadata)
330 -
331 def forward(
332 self,
333 input_ids: Optional[torch.LongTensor],
vllm/model_executor/models/prithvi_geospatial_mae.py
+8 -12
@@ -24,12 +24,13 @@ import torch.nn as nn
24 from transformers import BatchFeature
25
26 from vllm.config import VllmConfig
27 +from vllm.model_executor.layers.pooler import (AllPool, PoolerHead,
28 + PoolerIdentity, SimplePooler)
29 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30 from vllm.model_executor.models.interfaces import (IsAttentionFree,
31 SupportsMultiModal,
32 SupportsV0Only)
33 from vllm.model_executor.models.utils import AutoWeightsLoader
32 -from vllm.model_executor.pooling_metadata import PoolingMetadata
34 from vllm.multimodal import MULTIMODAL_REGISTRY
35 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
36 MultiModalInputs, MultiModalKwargs)
@@ -37,8 +38,7 @@ from vllm.multimodal.parse import MultiModalDataItems
38 from vllm.multimodal.processing import (BaseMultiModalProcessor,
39 BaseProcessingInfo, PromptUpdate)
40 from vllm.multimodal.profiling import BaseDummyInputsBuilder
40 -from vllm.sequence import (IntermediateTensors, PoolerOutput,
41 - PoolingSequenceGroupOutput)
41 +from vllm.sequence import IntermediateTensors
42
43
44 class PrithviGeoSpatialMAEProcessingInfo(BaseProcessingInfo):
@@ -116,7 +116,9 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
116 dummy_inputs=PrithviGeoSpatialMAEInputBuilder)
117 class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
118 SupportsV0Only):
119 - """ Prithvi Masked Autoencoder"""
119 + """Prithvi Masked Autoencoder"""
120 +
121 + is_pooling_model = True
122
123 @classmethod
124 def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
@@ -162,6 +164,8 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
164 "Only SemanticSegmentationTask is supported for now "
165 "by PrithviGeospatialMAE.")
166
167 + self.pooler = SimplePooler(AllPool(), PoolerHead(PoolerIdentity()))
168 +
169 def _parse_and_validate_multimodal_data(
170 self, **kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
171
@@ -189,7 +193,6 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
193 inputs_embeds: Optional[torch.Tensor] = None,
194 **kwargs: object,
195 ):
192 -
196 pixel_values, location_coords = (
197 self._parse_and_validate_multimodal_data(**kwargs))
198 model_output = self.model(pixel_values,
@@ -197,13 +200,6 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
200
201 return model_output.output
202
200 - def pooler(
201 - self,
202 - hidden_states: torch.Tensor,
203 - pooling_metadata: PoolingMetadata,
204 - ) -> Optional[PoolerOutput]:
205 - return PoolerOutput([PoolingSequenceGroupOutput(hidden_states)])
206 -
203 def load_weights(self, weights: Iterable[tuple[str,
204 torch.Tensor]]) -> set[str]:
205 params_list = []
vllm/model_executor/models/qwen2_rm.py
+9 -14
@@ -16,8 +16,7 @@ from vllm.config import VllmConfig
16 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17 RowParallelLinear)
18 from vllm.model_executor.layers.pooler import Pooler, PoolingType, SimplePooler
19 -from vllm.model_executor.pooling_metadata import PoolingMetadata
20 -from vllm.sequence import IntermediateTensors, PoolerOutput
19 +from vllm.sequence import IntermediateTensors
20
21 from .interfaces import SupportsLoRA, SupportsPP
22 from .qwen2 import Qwen2Model
@@ -25,6 +24,10 @@ from .utils import AutoWeightsLoader, maybe_prefix
24
25
26 class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
27 +
28 + is_pooling_model = True
29 + pooler: SimplePooler
30 +
31 packed_modules_mapping = {
32 "qkv_proj": [
33 "q_proj",
@@ -61,7 +64,6 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
64 quant_config=quant_config,
65 return_bias=False),
66 )
64 - self._pooler: SimplePooler
67 self.make_empty_intermediate_tensors = (
68 self.model.make_empty_intermediate_tensors)
69
@@ -80,13 +82,6 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
82 logits = self.score(hidden_states)
83 return logits
84
83 - def pooler(
84 - self,
85 - hidden_states: torch.Tensor,
86 - pooling_metadata: PoolingMetadata,
87 - ) -> Optional[PoolerOutput]:
88 - return self._pooler(hidden_states, pooling_metadata)
89 -
85 def load_weights(self, weights: Iterable[tuple[str,
86 torch.Tensor]]) -> set[str]:
87 loader = AutoWeightsLoader(self,
@@ -96,11 +91,11 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
91
92 class Qwen2ForRewardModel(Qwen2RewardBaseModel):
93
99 - def __init__(self, *, vllm_config, prefix=""):
94 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
95 vllm_config.model_config.hf_config.num_labels = 1
96 super().__init__(vllm_config=vllm_config, prefix=prefix)
97 pooler_config = vllm_config.model_config.pooler_config
103 - self._pooler = Pooler.from_config_with_defaults(
98 + self.pooler = Pooler.from_config_with_defaults(
99 pooler_config,
100 pooling_type=PoolingType.ALL,
101 normalize=False,
@@ -109,11 +104,11 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
104
105 class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
106
112 - def __init__(self, *, vllm_config, prefix=""):
107 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
108 vllm_config.model_config.hf_config.num_labels = 2
109 super().__init__(vllm_config=vllm_config, prefix=prefix)
110 pooler_config = vllm_config.model_config.pooler_config
116 - self._pooler = Pooler.from_config_with_defaults(
111 + self.pooler = Pooler.from_config_with_defaults(
112 pooler_config,
113 pooling_type=PoolingType.STEP,
114 normalize=False,
vllm/model_executor/models/roberta.py
+3 -10
@@ -15,8 +15,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
15 from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
16 from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
17 maybe_prefix)
18 -from vllm.model_executor.pooling_metadata import PoolingMetadata
19 -from vllm.sequence import IntermediateTensors, PoolerOutput
18 +from vllm.sequence import IntermediateTensors
19
20 from .bert_with_rope import BertWithRope, JinaRobertaModel
21 from .interfaces import SupportsCrossEncoding, SupportsV0Only
@@ -165,6 +164,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
164 _pooler: An instance of Pooler used for pooling operations.
165 """
166
167 + is_pooling_model = True
168 jina_to_vllm_mapper = WeightsMapper(
169 orig_to_new_substr={
170 'emb_ln': "embeddings.LayerNorm",
@@ -188,7 +188,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
188 add_pooling_layer=False)
189 self.classifier = RobertaClassificationHead(config)
190
191 - self._pooler = ClassifierPooler(
191 + self.pooler = ClassifierPooler(
192 vllm_config.model_config,
193 pooling=CLSPool(),
194 classifier=self.classifier,
@@ -198,13 +198,6 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
198 loader = AutoWeightsLoader(self)
199 return loader.load_weights(weights, mapper=self.jina_to_vllm_mapper)
200
201 - def pooler(
202 - self,
203 - hidden_states: torch.Tensor,
204 - pooling_metadata: PoolingMetadata,
205 - ) -> Optional[PoolerOutput]:
206 - return self._pooler(hidden_states, pooling_metadata)
207 -
201 def forward(
202 self,
203 input_ids: Optional[torch.Tensor],
vllm/pooling_params.py
+20 -11
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import TYPE_CHECKING, Any, Optional
4 +from typing import TYPE_CHECKING, Optional
5
6 import msgspec
7
@@ -15,24 +15,31 @@ class PoolingParams(
15 msgspec.Struct,
16 omit_defaults=True, # type: ignore[call-arg]
17 array_like=True): # type: ignore[call-arg]
18 - """API parameters for pooling models. This is currently a placeholder.
18 + """API parameters for pooling models. This
19
20 Attributes:
21 dimensions: Reduce the dimensions of embeddings
22 if model support matryoshka representation.
23 - additional_data: Any additional data needed for pooling.
23 """
24
25 dimensions: Optional[int] = None
26 +
27 use_cross_encoder: bool = False
28 - additional_data: Optional[Any] = None
28 + """Internal use only."""
29 +
30 + logits_processing_needs_token_ids: bool = False
31 + """Internal use only."""
32 +
33 output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
34
35 def clone(self) -> "PoolingParams":
36 """Returns a deep copy of the PoolingParams instance."""
33 - return PoolingParams(dimensions=self.dimensions,
34 - use_cross_encoder=self.use_cross_encoder,
35 - additional_data=self.additional_data)
37 + return PoolingParams(
38 + dimensions=self.dimensions,
39 + use_cross_encoder=self.use_cross_encoder,
40 + logits_processing_needs_token_ids=self.
41 + logits_processing_needs_token_ids,
42 + )
43
44 def verify(self, model_config: "ModelConfig") -> None:
45 if self.dimensions is not None:
@@ -54,10 +61,12 @@ class PoolingParams(
61 raise ValueError("Dimensions must be greater than 0")
62
63 def __repr__(self) -> str:
57 - return (f"PoolingParams("
58 - f"dimensions={self.dimensions}, "
59 - f"use_cross_encoder={self.use_cross_encoder}, "
60 - f"additional_metadata={self.additional_data})")
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 + )
70
71 def __post_init__(self) -> None:
72 assert self.output_kind == RequestOutputKind.FINAL_ONLY,\