@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c8ed39b9dd
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 10, 10:15 UTC

[Model] Reorganize pooling layers (#31973)

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

Cyrus Leung committed Jan 9, 2026 at 19:02 UTC c8ed39b9dd8efad4bed9547c52df7999c1735114
34 files changed +1290 -1143
.github/CODEOWNERS
+1 -1
@@ -153,7 +153,7 @@ mkdocs.yaml @hmellor
153 /vllm/entrypoints/pooling @noooop
154 /vllm/config/pooler.py @noooop
155 /vllm/pooling_params.py @noooop
156 -/vllm/model_executor/layers/pooler.py @noooop
156 +/vllm/model_executor/layers/pooler @noooop
157
158 # Security guide and policies
159 /docs/usage/security.md @russellb
tests/model_executor/test_model_load_with_params.py
+2 -1
@@ -5,7 +5,8 @@ import os
5
6 import pytest
7
8 -from vllm.model_executor.layers.pooler import CLSPool, DispatchPooler, MeanPool
8 +from vllm.model_executor.layers.pooler import DispatchPooler
9 +from vllm.model_executor.layers.pooler.seqwise import CLSPool, MeanPool
10 from vllm.model_executor.models.bert import BertEmbeddingModel
11 from vllm.model_executor.models.roberta import RobertaEmbeddingModel
12 from vllm.platforms import current_platform
tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py
+2 -7
@@ -7,7 +7,7 @@ import torch
7 import torch.nn as nn
8
9 from vllm.config import VllmConfig
10 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
10 +from vllm.model_executor.layers.pooler import DispatchPooler
11 from vllm.model_executor.models.gemma2 import Gemma2Model
12 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
13 from vllm.sequence import IntermediateTensors
@@ -28,12 +28,7 @@ class MyGemma2Embedding(nn.Module):
28 pooler_config = vllm_config.model_config.pooler_config
29 assert pooler_config is not None
30
31 - self.pooler = DispatchPooler(
32 - {
33 - "token_embed": Pooler.for_token_embed(pooler_config),
34 - "embed": Pooler.for_embed(pooler_config),
35 - }
36 - )
31 + self.pooler = DispatchPooler.for_embedding(pooler_config)
32
33 self.make_empty_intermediate_tensors = (
34 self.model.make_empty_intermediate_tensors
vllm/config/pooler.py
+4
@@ -88,6 +88,10 @@ class PoolerConfig:
88 # raise deprecated warning for softmax and activation
89 self.use_activation = get_use_activation(self)
90
91 + def get_pooling_type(self) -> PoolingTypeStr:
92 + assert self.pooling_type is not None, "Should be resolved by ModelConfig"
93 + return self.pooling_type
94 +
95 def compute_hash(self) -> str:
96 """
97 WARNING: Whenever a new field is added to this config,
vllm/model_executor/layers/pooler.py deleted
-845
@@ -1,845 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from abc import ABC, abstractmethod
4 -from collections.abc import Callable, Mapping, Set
5 -from dataclasses import dataclass
6 -from itertools import groupby
7 -from typing import TypeAlias, TypeVar
8 -
9 -import torch
10 -import torch.nn as nn
11 -import torch.nn.functional as F
12 -from transformers import PretrainedConfig
13 -
14 -from vllm.config import ModelConfig, get_current_vllm_config
15 -from vllm.config.pooler import PoolerConfig, PoolingTypeStr
16 -from vllm.logger import init_logger
17 -from vllm.model_executor.models.adapters import _load_st_projector
18 -from vllm.pooling_params import PoolingParams
19 -from vllm.tasks import PoolingTask
20 -from vllm.utils.import_utils import resolve_obj_by_qualname
21 -from vllm.v1.outputs import PoolerOutput, TokenPoolerOutput, TokenwisePoolerOutput
22 -from vllm.v1.pool.metadata import PoolingMetadata
23 -
24 -logger = init_logger(__name__)
25 -
26 -PoolingFn = Callable[
27 - [torch.Tensor | list[torch.Tensor], PoolingMetadata],
28 - torch.Tensor | list[torch.Tensor],
29 -]
30 -ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
31 -
32 -
33 -TokenPoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
34 -TokenwisePoolingMethodOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]
35 -TokenwisePoolingMethodOutputItem: TypeAlias = torch.Tensor | None
36 -PoolingMethodOutput: TypeAlias = TokenPoolingMethodOutput | TokenwisePoolingMethodOutput
37 -
38 -TokenPoolerHeadOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
39 -TokenwisePoolerHeadOutput: TypeAlias = torch.Tensor | None
40 -
41 -
42 -@dataclass(frozen=True)
43 -class ResolvedPoolingConfig:
44 - pooling_type: PoolingTypeStr
45 - task: PoolingTask
46 -
47 - @classmethod
48 - def from_config(
49 - cls,
50 - task: PoolingTask,
51 - pooler_config: PoolerConfig,
52 - ) -> "ResolvedPoolingConfig":
53 - assert pooler_config.pooling_type is not None
54 - return cls(task=task, pooling_type=pooler_config.pooling_type)
55 -
56 -
57 -@dataclass(frozen=True)
58 -class PoolingParamsUpdate:
59 - requires_token_ids: bool = False
60 - """Set this flag to enable `get_prompt_token_ids` for your pooler."""
61 -
62 - def apply(self, params: PoolingParams) -> None:
63 - params.requires_token_ids = self.requires_token_ids
64 -
65 -
66 -def get_classification_activation_function(config: PretrainedConfig):
67 - # Implement alignment with transformers ForSequenceClassificationLoss
68 - # https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92
69 - problem_type = getattr(config, "problem_type", "")
70 - if problem_type == "regression":
71 - return PoolerIdentity()
72 - if problem_type == "single_label_classification":
73 - return PoolerClassify()
74 - if problem_type == "multi_label_classification":
75 - return PoolerMultiLabelClassify()
76 - return PoolerClassify()
77 -
78 -
79 -def get_cross_encoder_activation_function(config: PretrainedConfig):
80 - function_name: str | None = None
81 - if (
82 - hasattr(config, "sentence_transformers")
83 - and "activation_fn" in config.sentence_transformers
84 - ):
85 - function_name = config.sentence_transformers["activation_fn"]
86 - elif (
87 - hasattr(config, "sbert_ce_default_activation_function")
88 - and config.sbert_ce_default_activation_function is not None
89 - ):
90 - function_name = config.sbert_ce_default_activation_function
91 -
92 - if function_name is not None:
93 - assert function_name.startswith("torch.nn.modules."), (
94 - "Loading of activation functions is restricted to "
95 - "torch.nn.modules for security reasons"
96 - )
97 - fn = resolve_obj_by_qualname(function_name)()
98 - return PoolerActivation.wraps(fn)
99 -
100 - return PoolerClassify()
101 -
102 -
103 -class PoolingMethod(nn.Module, ABC):
104 - @staticmethod
105 - def from_pooling_type(pooling_type: PoolingTypeStr) -> "PoolingMethod":
106 - if pooling_type == "LAST":
107 - return LastPool()
108 - if pooling_type == "ALL":
109 - return AllPool()
110 - if pooling_type == "CLS":
111 - return CLSPool()
112 - if pooling_type == "MEAN":
113 - return MeanPool()
114 - if pooling_type == "STEP":
115 - raise ValueError(
116 - "'STEP' pooling is handled by StepPooler "
117 - "and is not a standalone PoolingMethod."
118 - )
119 -
120 - raise NotImplementedError(f"Unsupported method: {pooling_type!r}")
121 -
122 - @abstractmethod
123 - def get_supported_tasks(self) -> Set[PoolingTask]:
124 - raise NotImplementedError
125 -
126 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
127 - return PoolingParamsUpdate()
128 -
129 - @abstractmethod
130 - def forward(
131 - self,
132 - hidden_states: torch.Tensor,
133 - pooling_metadata: PoolingMetadata,
134 - ) -> PoolingMethodOutput:
135 - raise NotImplementedError
136 -
137 -
138 -class CLSPool(PoolingMethod):
139 - def get_supported_tasks(self) -> Set[PoolingTask]:
140 - return {"token_embed", "token_classify", "embed", "classify", "score"}
141 -
142 - def forward(
143 - self,
144 - hidden_states: torch.Tensor,
145 - pooling_metadata: PoolingMetadata,
146 - ) -> TokenPoolingMethodOutput:
147 - pooling_cursor = pooling_metadata.get_pooling_cursor()
148 - assert not pooling_cursor.is_partial_prefill(), (
149 - "partial prefill not supported with CLS pooling"
150 - )
151 -
152 - return hidden_states[pooling_cursor.first_token_indices_gpu]
153 -
154 -
155 -class LastPool(PoolingMethod):
156 - def get_supported_tasks(self) -> Set[PoolingTask]:
157 - return {"token_embed", "token_classify", "embed", "classify", "score"}
158 -
159 - def forward(
160 - self,
161 - hidden_states: torch.Tensor,
162 - pooling_metadata: PoolingMetadata,
163 - ) -> TokenPoolingMethodOutput:
164 - pooling_cursor = pooling_metadata.get_pooling_cursor()
165 - return hidden_states[pooling_cursor.last_token_indices_gpu]
166 -
167 -
168 -class AllPool(PoolingMethod):
169 - def __init__(self):
170 - super().__init__()
171 -
172 - vllm_config = get_current_vllm_config()
173 - self.enable_chunked_prefill = (
174 - vllm_config.scheduler_config.enable_chunked_prefill
175 - )
176 -
177 - def get_supported_tasks(self) -> Set[PoolingTask]:
178 - return {"token_embed", "token_classify"}
179 -
180 - def forward(
181 - self,
182 - hidden_states: torch.Tensor,
183 - pooling_metadata: PoolingMetadata,
184 - ) -> TokenwisePoolingMethodOutput:
185 - pooling_cursor = pooling_metadata.get_pooling_cursor()
186 - hidden_states_all = hidden_states.split(
187 - pooling_cursor.num_scheduled_tokens_cpu.tolist()
188 - )
189 - hidden_states_lst = [hidden_states_all[i] for i in pooling_cursor.index]
190 -
191 - if not self.enable_chunked_prefill:
192 - return hidden_states_lst
193 -
194 - pooling_states = pooling_metadata.pooling_states
195 -
196 - # If chunked_prefill is enabled
197 - # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache
198 - for p, hs_chunk in zip(pooling_states, hidden_states_lst):
199 - p.hidden_states_cache.append(hs_chunk)
200 -
201 - # 2. Once prefill is finished, send hidden_states_cache to PoolerHead
202 - output_list = list[torch.Tensor | None]()
203 - for p, finished in zip(pooling_states, pooling_cursor.is_finished()):
204 - if finished:
205 - hidden_states_cache = p.hidden_states_cache
206 - if len(hidden_states_cache) == 1:
207 - output_list.append(hidden_states_cache[0])
208 - else:
209 - output_list.append(torch.concat(hidden_states_cache, dim=0))
210 - p.clean()
211 - else:
212 - output_list.append(None)
213 -
214 - return output_list
215 -
216 -
217 -class MeanPool(PoolingMethod):
218 - def get_supported_tasks(self) -> Set[PoolingTask]:
219 - return {"token_embed", "token_classify", "embed", "classify", "score"}
220 -
221 - def forward(
222 - self,
223 - hidden_states: torch.Tensor,
224 - pooling_metadata: PoolingMetadata,
225 - ) -> TokenPoolingMethodOutput:
226 - pooling_cursor = pooling_metadata.get_pooling_cursor()
227 - assert not pooling_cursor.is_partial_prefill(), (
228 - "partial prefill not supported with MEAN pooling"
229 - )
230 -
231 - prompt_lens = pooling_cursor.prompt_lens_cpu.to(
232 - hidden_states.device, non_blocking=True
233 - )
234 -
235 - # Use float32 for torch.cumsum in MeanPool,
236 - # otherwise precision will be lost significantly.
237 - cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
238 -
239 - start_indices = pooling_cursor.first_token_indices_gpu
240 - end_indices = pooling_cursor.last_token_indices_gpu
241 - return (
242 - cumsum[end_indices] - cumsum[start_indices] + hidden_states[start_indices]
243 - ) / prompt_lens.unsqueeze(1)
244 -
245 -
246 -_T = TypeVar("_T", torch.Tensor, list[torch.Tensor])
247 -
248 -
249 -class BasePoolerActivation(nn.Module, ABC):
250 - @abstractmethod
251 - def forward(self, pooled_data: _T) -> _T:
252 - # shape:
253 - # classify (& score) -> (batch_size, num_classes)
254 - # embed -> (batch_size, embedding_dim) or list(embedding_dim)
255 - # (batch_size, dimensions) or list(dimensions) if using MRL
256 - raise NotImplementedError
257 -
258 -
259 -class PoolerActivation(BasePoolerActivation):
260 - @staticmethod
261 - def wraps(module: nn.Module):
262 - if isinstance(module, nn.Identity):
263 - return PoolerIdentity()
264 - if isinstance(module, (nn.Sigmoid, nn.Softmax)):
265 - return PoolerClassify()
266 -
267 - return LambdaPoolerActivation(module)
268 -
269 - @abstractmethod
270 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
271 - raise NotImplementedError
272 -
273 - def forward(self, pooled_data: _T) -> _T:
274 - if isinstance(pooled_data, list):
275 - return [self.forward_chunk(data) for data in pooled_data]
276 -
277 - return self.forward_chunk(pooled_data)
278 -
279 -
280 -class PoolerIdentity(PoolerActivation):
281 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
282 - return pooled_data
283 -
284 -
285 -class PoolerNormalize(PoolerActivation):
286 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
287 - return F.normalize(pooled_data, p=2, dim=-1)
288 -
289 -
290 -class PoolerMultiLabelClassify(PoolerActivation):
291 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
292 - return F.sigmoid(pooled_data)
293 -
294 -
295 -class PoolerClassify(PoolerActivation):
296 - def __init__(self, *, static_num_labels: bool = True) -> None:
297 - super().__init__()
298 -
299 - if static_num_labels:
300 - vllm_config = get_current_vllm_config()
301 - self.num_labels = getattr(
302 - vllm_config.model_config.hf_config, "num_labels", 0
303 - )
304 - if self.num_labels == 0:
305 - logger.warning(
306 - "num_labels should be > 0 for classification"
307 - "models, falling back to softmax. "
308 - "Please check if the configuration is correct."
309 - )
310 - else:
311 - self.num_labels = None
312 -
313 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
314 - num_labels = (
315 - self.num_labels if self.num_labels is not None else pooled_data.shape[-1]
316 - )
317 -
318 - if num_labels < 2:
319 - return F.sigmoid(pooled_data)
320 -
321 - return F.softmax(pooled_data, dim=-1)
322 -
323 -
324 -class LambdaPoolerActivation(PoolerActivation):
325 - def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
326 - super().__init__()
327 -
328 - self.fn = fn
329 -
330 - def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
331 - return self.fn(pooled_data)
332 -
333 -
334 -class Pooler(nn.Module, ABC):
335 - """The interface required for all poolers used in pooling models in vLLM."""
336 -
337 - @staticmethod
338 - def for_token_embed(pooler_config: PoolerConfig):
339 - head = TokenEmbeddingPoolerHead()
340 -
341 - if pooler_config.pooling_type == "STEP":
342 - return StepPooler(head=head)
343 -
344 - return AllPooler(head=head)
345 -
346 - @staticmethod
347 - def for_token_classify(
348 - pooler_config: PoolerConfig,
349 - classifier: ClassifierFn | None = None,
350 - act_fn: PoolerActivation | str | None = None,
351 - ):
352 - head = TokenClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
353 -
354 - if pooler_config.pooling_type == "STEP":
355 - return StepPooler(head=head)
356 -
357 - return AllPooler(head=head)
358 -
359 - @staticmethod
360 - def for_embed(pooler_config: PoolerConfig):
361 - resolved_config = ResolvedPoolingConfig.from_config(
362 - task="embed",
363 - pooler_config=pooler_config,
364 - )
365 -
366 - pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
367 - head = EmbeddingPoolerHead()
368 -
369 - return SimplePooler(pooling=pooling, head=head)
370 -
371 - @staticmethod
372 - def for_classify(
373 - pooler_config: PoolerConfig,
374 - classifier: ClassifierFn | None,
375 - act_fn: PoolerActivation | str | None = None,
376 - ):
377 - resolved_config = ResolvedPoolingConfig.from_config(
378 - task="classify",
379 - pooler_config=pooler_config,
380 - )
381 -
382 - pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
383 -
384 - return ClassifierPooler(
385 - pooling=pooling,
386 - classifier=classifier,
387 - act_fn=act_fn,
388 - )
389 -
390 - @abstractmethod
391 - def get_supported_tasks(self) -> Set[PoolingTask]:
392 - """Determine which pooling tasks are supported."""
393 - raise NotImplementedError
394 -
395 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
396 - """
397 - Construct the updated pooling parameters to use for a supported task.
398 - """
399 - return PoolingParamsUpdate()
400 -
401 - @abstractmethod
402 - def forward(
403 - self,
404 - hidden_states: torch.Tensor,
405 - pooling_metadata: PoolingMetadata,
406 - ) -> PoolerOutput:
407 - raise NotImplementedError
408 -
409 -
410 -class DummyPooler(Pooler):
411 - def get_supported_tasks(self) -> Set[PoolingTask]:
412 - return {"plugin", "score"}
413 -
414 - def forward(
415 - self,
416 - hidden_states: torch.Tensor,
417 - pooling_metadata: PoolingMetadata,
418 - ) -> PoolerOutput:
419 - return hidden_states
420 -
421 -
422 -class TokenPoolerHead(nn.Module, ABC):
423 - """Applicable to pooling strategies that output one token."""
424 -
425 - @abstractmethod
426 - def forward(
427 - self,
428 - pooled_data: TokenPoolingMethodOutput,
429 - pooling_metadata: PoolingMetadata,
430 - ) -> TokenPoolerHeadOutput:
431 - raise NotImplementedError
432 -
433 -
434 -class EmbeddingPoolerHead(TokenPoolerHead):
435 - def __init__(self) -> None:
436 - super().__init__()
437 -
438 - # Load ST projector if available
439 - vllm_config = get_current_vllm_config()
440 - self.projector = (
441 - _load_st_projector(vllm_config.model_config) if vllm_config else None
442 - )
443 - self.head_dtype = vllm_config.model_config.head_dtype
444 -
445 - self.activation = PoolerNormalize()
446 -
447 - def forward(
448 - self,
449 - pooled_data: TokenPoolingMethodOutput,
450 - pooling_metadata: PoolingMetadata,
451 - ) -> TokenPoolerHeadOutput:
452 - if isinstance(pooled_data, list):
453 - pooled_data = torch.stack(pooled_data)
454 - # pooled_data shape: [batchsize, hidden_dimension]
455 -
456 - pooled_data = pooled_data.to(self.head_dtype)
457 -
458 - # Apply ST projector
459 - if self.projector is not None:
460 - pooled_data = self.projector(pooled_data)
461 - # pooled_data shape: [batchsize, embedding_dimension]
462 -
463 - pooling_params = pooling_metadata.pooling_params
464 -
465 - # for matryoshka representation
466 - dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params]
467 - if any(d is not None for d in dimensions_list):
468 - # change the output dimension
469 - assert len(pooled_data) == len(dimensions_list)
470 - if len(set(dimensions_list)) == 1 and not isinstance(pooled_data, list):
471 - # if all dimensions are the same
472 - d = dimensions_list[0]
473 - pooled_data = pooled_data[..., :d]
474 - else:
475 - pooled_data = [
476 - vecs if d is None else vecs[..., :d]
477 - for vecs, d in zip(pooled_data, dimensions_list)
478 - ]
479 -
480 - # for normalize
481 - flags = [p.normalize for p in pooling_params]
482 - if len(set(flags)) == 1:
483 - if flags[0]:
484 - pooled_data = self.activation(pooled_data)
485 - else:
486 - pooled_data = [
487 - self.activation(vecs) if f else vecs
488 - for vecs, f in zip(pooled_data, flags)
489 - ]
490 -
491 - # pooled_data shape: [batchsize, embedding_dimension]
492 - return pooled_data
493 -
494 -
495 -class SimplePooler(Pooler):
496 - """A layer that pools specific information from hidden states.
497 -
498 - This layer does the following:
499 - 1. Extracts specific tokens or aggregates data based on pooling method.
500 - 2. Normalizes output if specified.
501 - 3. Returns structured results as `PoolerOutput`.
502 - """
503 -
504 - def __init__(self, pooling: PoolingMethod, head: TokenPoolerHead) -> None:
505 - super().__init__()
506 -
507 - self.pooling = pooling
508 - self.head = head
509 -
510 - def get_supported_tasks(self) -> Set[PoolingTask]:
511 - return self.pooling.get_supported_tasks()
512 -
513 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
514 - return self.pooling.get_pooling_updates(task)
515 -
516 - def forward(
517 - self,
518 - hidden_states: torch.Tensor,
519 - pooling_metadata: PoolingMetadata,
520 - ) -> TokenPoolerHeadOutput:
521 - pooled_data = self.pooling(hidden_states, pooling_metadata)
522 - pooled_data = self.head(pooled_data, pooling_metadata)
523 - return pooled_data
524 -
525 -
526 -class ClassifierPooler(Pooler):
527 - """A pooling layer for classification tasks.
528 -
529 - This layer does the following:
530 - 1. Applies a classification layer to the hidden states.
531 - 2. Optionally applies a pooler layer.
532 - 3. Applies an activation function to the output.
533 - """
534 -
535 - @staticmethod
536 - def act_fn_for_seq_cls(model_config: ModelConfig):
537 - return get_classification_activation_function(model_config.hf_config)
538 -
539 - @staticmethod
540 - def act_fn_for_cross_encoder(model_config: ModelConfig):
541 - return get_cross_encoder_activation_function(model_config.hf_config)
542 -
543 - @staticmethod
544 - def resolve_act_fn(
545 - model_config: ModelConfig,
546 - static_num_labels: bool = True,
547 - act_fn: PoolerActivation | str | None = None,
548 - ):
549 - if isinstance(act_fn, str):
550 - if act_fn == "classify":
551 - return ClassifierPooler.act_fn_for_seq_cls(model_config)
552 - elif act_fn == "score":
553 - return ClassifierPooler.act_fn_for_cross_encoder(model_config)
554 - else:
555 - raise ValueError(f"act_fn [{act_fn=}] not supported.")
556 - elif act_fn is None:
557 - return PoolerClassify(static_num_labels=static_num_labels)
558 - else:
559 - assert callable(act_fn)
560 - return act_fn
561 -
562 - def __init__(
563 - self,
564 - pooling: PoolingFn,
565 - classifier: ClassifierFn | None,
566 - act_fn: PoolerActivation | str | None = None,
567 - ) -> None:
568 - super().__init__()
569 -
570 - vllm_config = get_current_vllm_config()
571 - self.pooling = pooling
572 - self.classifier = classifier
573 - self.act_fn = self.resolve_act_fn(
574 - vllm_config.model_config, static_num_labels=True, act_fn=act_fn
575 - )
576 - self.logit_bias: float | None = (
577 - vllm_config.model_config.pooler_config.logit_bias
578 - )
579 - self.head_dtype = vllm_config.model_config.head_dtype
580 -
581 - def get_supported_tasks(self) -> Set[PoolingTask]:
582 - return {"classify", "score"}
583 -
584 - def forward(
585 - self,
586 - hidden_states: torch.Tensor,
587 - pooling_metadata: PoolingMetadata,
588 - ) -> TokenPoolerOutput:
589 - pooled_data = self.pooling(hidden_states, pooling_metadata)
590 - if isinstance(pooled_data, list):
591 - pooled_data = torch.stack(pooled_data)
592 - # pooled_data shape: [batchsize, hidden_size]
593 -
594 - pooled_data = pooled_data.to(self.head_dtype)
595 -
596 - if self.classifier is not None:
597 - pooled_data = self.classifier(pooled_data)
598 - # pooled_data shape: [batchsize, num_labels]
599 -
600 - if self.logit_bias is not None:
601 - pooled_data -= self.logit_bias
602 -
603 - pooling_params = pooling_metadata.pooling_params
604 - flags = [p.use_activation for p in pooling_params]
605 -
606 - if len(set(flags)) == 1:
607 - scores = self.act_fn(pooled_data) if flags[0] else pooled_data
608 - else:
609 - scores = [
610 - self.act_fn(vecs) if f else vecs for vecs, f in zip(pooled_data, flags)
611 - ]
612 -
613 - # scores shape: [batchsize, num_labels]
614 - return scores
615 -
616 -
617 -class TokenwisePoolerHead(nn.Module, ABC):
618 - """Applicable to pooling strategies that output multiple tokens."""
619 -
620 - @abstractmethod
621 - def forward(
622 - self,
623 - pooled_data: TokenwisePoolingMethodOutputItem,
624 - pooling_param: PoolingParams,
625 - ) -> TokenwisePoolerHeadOutput:
626 - raise NotImplementedError
627 -
628 -
629 -class TokenEmbeddingPoolerHead(TokenwisePoolerHead):
630 - def __init__(self) -> None:
631 - super().__init__()
632 -
633 - # Load ST projector if available
634 - vllm_config = get_current_vllm_config()
635 - self.projector = (
636 - _load_st_projector(vllm_config.model_config) if vllm_config else None
637 - )
638 - self.head_dtype = vllm_config.model_config.head_dtype
639 -
640 - self.activation = PoolerNormalize()
641 -
642 - def forward(
643 - self,
644 - pooled_data: TokenwisePoolingMethodOutputItem,
645 - pooling_param: PoolingParams,
646 - ) -> TokenwisePoolerHeadOutput:
647 - # for unfinished chunked prefill
648 - if pooled_data is None:
649 - return None
650 -
651 - pooled_data = pooled_data.to(self.head_dtype)
652 - # pooled_data shape: [n_tokens, hidden_dimension]
653 -
654 - # Apply ST projector
655 - if self.projector is not None:
656 - pooled_data = self.projector(pooled_data)
657 - # pooled_data shape: [n_tokens, embedding_dimension]
658 -
659 - # for matryoshka representation
660 - pooled_data = pooled_data[..., : pooling_param.dimensions]
661 -
662 - # for normalize
663 - if pooling_param.normalize:
664 - pooled_data = self.activation(pooled_data)
665 -
666 - # pooled_data shape: [n_tokens, embedding_dimension]
667 - return pooled_data
668 -
669 -
670 -class TokenClassifierPoolerHead(TokenwisePoolerHead):
671 - def __init__(
672 - self,
673 - classifier: ClassifierFn | None,
674 - act_fn: PoolerActivation | str | None = None,
675 - ) -> None:
676 - super().__init__()
677 -
678 - vllm_config = get_current_vllm_config()
679 -
680 - self.classifier = classifier
681 - self.logit_bias: float | None = (
682 - vllm_config.model_config.pooler_config.logit_bias
683 - )
684 - self.head_dtype = vllm_config.model_config.head_dtype
685 -
686 - self.activation = ClassifierPooler.resolve_act_fn(
687 - vllm_config.model_config, static_num_labels=False, act_fn=act_fn
688 - )
689 -
690 - def forward(
691 - self,
692 - pooled_data: TokenwisePoolingMethodOutputItem,
693 - pooling_param: PoolingParams,
694 - ) -> TokenwisePoolerHeadOutput:
695 - # for unfinished chunked prefill
696 - if pooled_data is None:
697 - return None
698 -
699 - pooled_data = pooled_data.to(self.head_dtype)
700 - # hidden_states shape: [n_token, hidden_size]
701 -
702 - if self.classifier is not None:
703 - scores = self.classifier(pooled_data)
704 - else:
705 - scores = pooled_data
706 - # scores shape: [n_token, num_labels]
707 -
708 - if self.logit_bias is not None:
709 - scores -= self.logit_bias
710 -
711 - if pooling_param.use_activation:
712 - scores = self.activation(scores)
713 -
714 - # scores shape: [n_token, num_labels]
715 - return scores
716 -
717 -
718 -class AllPooler(Pooler):
719 - def __init__(self, head: TokenwisePoolerHead) -> None:
720 - super().__init__()
721 -
722 - self.pooling = AllPool()
723 - self.head = head
724 -
725 - def get_supported_tasks(self) -> Set[PoolingTask]:
726 - return {"token_embed", "token_classify"}
727 -
728 - def forward(
729 - self,
730 - hidden_states: torch.Tensor,
731 - pooling_metadata: PoolingMetadata,
732 - ) -> TokenwisePoolerOutput:
733 - pooled_data = self.pooling(hidden_states, pooling_metadata)
734 - pooling_params = pooling_metadata.pooling_params
735 - assert len(pooled_data) == len(pooling_params)
736 -
737 - return [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
738 -
739 -
740 -class StepPooler(Pooler):
741 - def __init__(self, head: TokenwisePoolerHead) -> None:
742 - super().__init__()
743 -
744 - self.pooling = AllPool()
745 - self.head = head
746 -
747 - def extract_states(
748 - self,
749 - hidden_states: torch.Tensor,
750 - pooling_metadata: PoolingMetadata,
751 - ) -> list[torch.Tensor | None]:
752 - pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
753 - prompt_token_ids = pooling_metadata.get_prompt_token_ids()
754 - pooling_params = pooling_metadata.pooling_params
755 -
756 - pooled_data = list[torch.Tensor | None]()
757 - for data, token_id, pooling_param in zip(
758 - pooled_data_lst, prompt_token_ids, pooling_params
759 - ):
760 - # for unfinished chunked prefill
761 - if data is None:
762 - pooled_data.append(data)
763 - continue
764 -
765 - step_tag_id = pooling_param.step_tag_id
766 - returned_token_ids = pooling_param.returned_token_ids
767 -
768 - if returned_token_ids is not None and len(returned_token_ids) > 0:
769 - data = data[:, returned_token_ids]
770 -
771 - if step_tag_id is not None:
772 - data = data[token_id == step_tag_id]
773 - pooled_data.append(data)
774 -
775 - return pooled_data
776 -
777 - def get_supported_tasks(self) -> Set[PoolingTask]:
778 - return {"token_embed", "token_classify"}
779 -
780 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
781 - return PoolingParamsUpdate(requires_token_ids=True)
782 -
783 - def forward(
784 - self,
785 - hidden_states: torch.Tensor,
786 - pooling_metadata: PoolingMetadata,
787 - ) -> TokenwisePoolerOutput:
788 - pooled_data = self.extract_states(hidden_states, pooling_metadata)
789 - pooling_params = pooling_metadata.pooling_params
790 - assert len(pooled_data) == len(pooling_params)
791 -
792 - return [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
793 -
794 -
795 -class DispatchPooler(Pooler):
796 - """Dispatches calls to a sub-pooler based on the pooling task."""
797 -
798 - def __init__(self, poolers_by_task: Mapping[PoolingTask, Pooler]) -> None:
799 - super().__init__()
800 -
801 - for task, pooler in poolers_by_task.items():
802 - if task not in pooler.get_supported_tasks():
803 - raise ValueError(
804 - f"{pooler=} does not support {task=}. "
805 - f"Supported tasks: {pooler.get_supported_tasks()}"
806 - )
807 -
808 - self.poolers_by_task = poolers_by_task
809 -
810 - def get_supported_tasks(self) -> Set[PoolingTask]:
811 - return set(self.poolers_by_task)
812 -
813 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
814 - return self.poolers_by_task[task].get_pooling_updates(task)
815 -
816 - def forward(
817 - self,
818 - hidden_states: torch.Tensor,
819 - pooling_metadata: PoolingMetadata,
820 - ) -> PoolerOutput:
821 - poolers_by_task = self.poolers_by_task
822 -
823 - outputs = list[torch.Tensor | None]()
824 - offset = 0
825 - for task, group in groupby(pooling_metadata.tasks):
826 - if not (pooler := poolers_by_task.get(task)):
827 - raise ValueError(
828 - f"Unsupported task: {task} "
829 - f"Supported tasks: {self.get_supported_tasks()}"
830 - )
831 -
832 - num_items = len(list(group))
833 - group_output: PoolerOutput = pooler(
834 - hidden_states,
835 - pooling_metadata[offset : offset + num_items],
836 - )
837 -
838 - outputs.extend(group_output)
839 - offset += num_items
840 -
841 - return outputs
842 -
843 - def extra_repr(self) -> str:
844 - s = f"supported_task={self.get_supported_tasks()}"
845 - return s
vllm/model_executor/layers/pooler/__init__.py new
+5
@@ -0,0 +1,5 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from .abstract import *
4 +from .common import *
5 +from .special import *
vllm/model_executor/layers/pooler/abstract.py new
+39
@@ -0,0 +1,39 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Set
5 +
6 +import torch
7 +import torch.nn as nn
8 +
9 +from vllm.tasks import PoolingTask
10 +from vllm.v1.outputs import PoolerOutput
11 +from vllm.v1.pool.metadata import PoolingMetadata
12 +
13 +from .common import PoolingParamsUpdate
14 +
15 +
16 +class Pooler(nn.Module, ABC):
17 + """The interface required for all poolers used in pooling models in vLLM."""
18 +
19 + @abstractmethod
20 + def get_supported_tasks(self) -> Set[PoolingTask]:
21 + """Determine which pooling tasks are supported."""
22 + raise NotImplementedError
23 +
24 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
25 + """
26 + Construct the updated pooling parameters to use for a supported task.
27 + """
28 + return PoolingParamsUpdate()
29 +
30 + @abstractmethod
31 + def forward(
32 + self,
33 + hidden_states: torch.Tensor,
34 + pooling_metadata: PoolingMetadata,
35 + ) -> PoolerOutput:
36 + raise NotImplementedError
37 +
38 +
39 +__all__ = ["Pooler"]
vllm/model_executor/layers/pooler/activations.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Callable
5 +from typing import TypeVar
6 +
7 +import torch
8 +import torch.nn as nn
9 +import torch.nn.functional as F
10 +from transformers import PretrainedConfig
11 +
12 +from vllm.config import ModelConfig, get_current_vllm_config
13 +from vllm.logger import init_logger
14 +from vllm.utils.import_utils import resolve_obj_by_qualname
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +def get_classification_act_fn(
20 + config: PretrainedConfig,
21 +) -> "PoolerActivation":
22 + # Implement alignment with transformers ForSequenceClassificationLoss
23 + # https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92
24 + problem_type = getattr(config, "problem_type", "")
25 + if problem_type == "regression":
26 + return PoolerIdentity()
27 + if problem_type == "single_label_classification":
28 + return PoolerClassify()
29 + if problem_type == "multi_label_classification":
30 + return PoolerMultiLabelClassify()
31 +
32 + return PoolerClassify()
33 +
34 +
35 +def get_cross_encoder_act_fn(
36 + config: PretrainedConfig,
37 +) -> "PoolerActivation":
38 + function_name: str | None = None
39 + if (
40 + hasattr(config, "sentence_transformers")
41 + and "activation_fn" in config.sentence_transformers
42 + ):
43 + function_name = config.sentence_transformers["activation_fn"]
44 + elif (
45 + hasattr(config, "sbert_ce_default_activation_function")
46 + and config.sbert_ce_default_activation_function is not None
47 + ):
48 + function_name = config.sbert_ce_default_activation_function
49 +
50 + if function_name is not None:
51 + assert function_name.startswith("torch.nn.modules."), (
52 + "Loading of activation functions is restricted to "
53 + "torch.nn.modules for security reasons"
54 + )
55 + fn = resolve_obj_by_qualname(function_name)()
56 + return PoolerActivation.wraps(fn)
57 +
58 + return PoolerClassify()
59 +
60 +
61 +def resolve_classifier_act_fn(
62 + model_config: ModelConfig,
63 + static_num_labels: bool = True,
64 + act_fn: "PoolerActivation | str | None" = None,
65 +):
66 + if isinstance(act_fn, str):
67 + if act_fn == "classify":
68 + return get_classification_act_fn(model_config.hf_config)
69 + if act_fn == "score":
70 + return get_cross_encoder_act_fn(model_config.hf_config)
71 +
72 + raise ValueError(f"act_fn [{act_fn=}] not supported.")
73 +
74 + if act_fn is None:
75 + return PoolerClassify(static_num_labels=static_num_labels)
76 +
77 + assert callable(act_fn)
78 + return act_fn
79 +
80 +
81 +_T = TypeVar("_T", torch.Tensor, list[torch.Tensor])
82 +
83 +
84 +class PoolerActivation(nn.Module, ABC):
85 + @staticmethod
86 + def wraps(module: nn.Module):
87 + if isinstance(module, nn.Identity):
88 + return PoolerIdentity()
89 + if isinstance(module, (nn.Sigmoid, nn.Softmax)):
90 + return PoolerClassify()
91 +
92 + return LambdaPoolerActivation(module)
93 +
94 + @abstractmethod
95 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
96 + raise NotImplementedError
97 +
98 + def forward(self, pooled_data: _T) -> _T:
99 + # shape:
100 + # classify (& score) -> (batch_size, num_classes)
101 + # embed -> (batch_size, embedding_dim) or list(embedding_dim)
102 + # (batch_size, dimensions) or list(dimensions) if using MRL
103 + if isinstance(pooled_data, list):
104 + return [self.forward_chunk(data) for data in pooled_data]
105 +
106 + return self.forward_chunk(pooled_data)
107 +
108 +
109 +class PoolerIdentity(PoolerActivation):
110 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
111 + return pooled_data
112 +
113 +
114 +class PoolerNormalize(PoolerActivation):
115 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
116 + return F.normalize(pooled_data, p=2, dim=-1)
117 +
118 +
119 +class PoolerMultiLabelClassify(PoolerActivation):
120 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
121 + return F.sigmoid(pooled_data)
122 +
123 +
124 +class PoolerClassify(PoolerActivation):
125 + def __init__(self, *, static_num_labels: bool = True) -> None:
126 + super().__init__()
127 +
128 + if static_num_labels:
129 + vllm_config = get_current_vllm_config()
130 + model_config = vllm_config.model_config
131 + num_labels = getattr(model_config.hf_config, "num_labels", 0)
132 + else:
133 + num_labels = None
134 +
135 + if num_labels == 0:
136 + logger.warning(
137 + "num_labels should be > 0 for classification "
138 + "models, falling back to softmax. "
139 + "Please check if the configuration is correct."
140 + )
141 +
142 + self.num_labels = num_labels
143 +
144 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
145 + num_labels = self.num_labels
146 + if num_labels is None:
147 + num_labels = pooled_data.shape[-1]
148 +
149 + if num_labels < 2:
150 + return F.sigmoid(pooled_data)
151 +
152 + return F.softmax(pooled_data, dim=-1)
153 +
154 +
155 +class LambdaPoolerActivation(PoolerActivation):
156 + def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
157 + super().__init__()
158 +
159 + self.fn = fn
160 +
161 + def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
162 + return self.fn(pooled_data)
vllm/model_executor/layers/pooler/common.py new
+27
@@ -0,0 +1,27 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +from dataclasses import dataclass
5 +
6 +import torch
7 +
8 +from vllm.pooling_params import PoolingParams
9 +
10 +ClassifierFn = Callable[[torch.Tensor], torch.Tensor]
11 +
12 +
13 +@dataclass(frozen=True)
14 +class PoolingParamsUpdate:
15 + requires_token_ids: bool = False
16 + """Set this flag to enable `get_prompt_token_ids` for your pooler."""
17 +
18 + def __or__(self, other: "PoolingParamsUpdate") -> "PoolingParamsUpdate":
19 + return PoolingParamsUpdate(
20 + requires_token_ids=self.requires_token_ids or other.requires_token_ids,
21 + )
22 +
23 + def apply(self, params: PoolingParams) -> None:
24 + params.requires_token_ids = self.requires_token_ids
25 +
26 +
27 +__all__ = ["ClassifierFn", "PoolingParamsUpdate"]
vllm/model_executor/layers/pooler/seqwise/__init__.py new
+45
@@ -0,0 +1,45 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Poolers that produce an output aggregating all tokens in the sequence."""
4 +
5 +from .heads import (
6 + ClassifierPoolerHead,
7 + EmbeddingPoolerHead,
8 + SequencePoolerHead,
9 + SequencePoolerHeadOutput,
10 +)
11 +from .methods import (
12 + CLSPool,
13 + LastPool,
14 + MeanPool,
15 + SequencePoolingMethod,
16 + SequencePoolingMethodOutput,
17 + get_seq_pooling_method,
18 +)
19 +from .poolers import (
20 + SequencePooler,
21 + SequencePoolerOutput,
22 + SequencePoolingFn,
23 + SequencePoolingHeadFn,
24 + pooler_for_classify,
25 + pooler_for_embed,
26 +)
27 +
28 +__all__ = [
29 + "SequencePoolerHead",
30 + "SequencePoolerHeadOutput",
31 + "ClassifierPoolerHead",
32 + "EmbeddingPoolerHead",
33 + "SequencePoolingMethod",
34 + "SequencePoolingMethodOutput",
35 + "CLSPool",
36 + "LastPool",
37 + "MeanPool",
38 + "get_seq_pooling_method",
39 + "SequencePooler",
40 + "SequencePoolingFn",
41 + "SequencePoolingHeadFn",
42 + "SequencePoolerOutput",
43 + "pooler_for_classify",
44 + "pooler_for_embed",
45 +]
vllm/model_executor/layers/pooler/seqwise/heads.py new
+157
@@ -0,0 +1,157 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Set
5 +from typing import TypeAlias
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import get_current_vllm_config
11 +from vllm.model_executor.layers.pooler import ClassifierFn
12 +from vllm.model_executor.layers.pooler.activations import (
13 + PoolerActivation,
14 + PoolerNormalize,
15 + resolve_classifier_act_fn,
16 +)
17 +from vllm.model_executor.models.adapters import _load_st_projector
18 +from vllm.tasks import PoolingTask
19 +from vllm.v1.pool.metadata import PoolingMetadata
20 +
21 +from .methods import SequencePoolingMethodOutput
22 +
23 +SequencePoolerHeadOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
24 +
25 +
26 +class SequencePoolerHead(nn.Module, ABC):
27 + @abstractmethod
28 + def get_supported_tasks(self) -> Set[PoolingTask]:
29 + raise NotImplementedError
30 +
31 + @abstractmethod
32 + def forward(
33 + self,
34 + pooled_data: SequencePoolingMethodOutput,
35 + pooling_metadata: PoolingMetadata,
36 + ) -> SequencePoolerHeadOutput:
37 + raise NotImplementedError
38 +
39 +
40 +class EmbeddingPoolerHead(SequencePoolerHead):
41 + def __init__(self) -> None:
42 + super().__init__()
43 +
44 + # Load ST projector if available
45 + vllm_config = get_current_vllm_config()
46 + model_config = vllm_config.model_config
47 +
48 + self.projector = _load_st_projector(model_config)
49 + self.head_dtype = model_config.head_dtype
50 +
51 + self.activation = PoolerNormalize()
52 +
53 + def get_supported_tasks(self) -> Set[PoolingTask]:
54 + return {"embed"}
55 +
56 + def forward(
57 + self,
58 + pooled_data: SequencePoolingMethodOutput,
59 + pooling_metadata: PoolingMetadata,
60 + ) -> SequencePoolerHeadOutput:
61 + pooling_params = pooling_metadata.pooling_params
62 + assert len(pooled_data) == len(pooling_params)
63 +
64 + if isinstance(pooled_data, list):
65 + pooled_data = torch.stack(pooled_data)
66 + # pooled_data shape: [batchsize, hidden_dimension]
67 +
68 + pooled_data = pooled_data.to(self.head_dtype)
69 +
70 + # Apply ST projector
71 + if self.projector is not None:
72 + pooled_data = self.projector(pooled_data)
73 + # pooled_data shape: [batchsize, embedding_dimension]
74 +
75 + # for matryoshka representation
76 + dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params]
77 + if any(d is not None for d in dimensions_list):
78 + # change the output dimension
79 + assert len(pooled_data) == len(dimensions_list)
80 + if len(set(dimensions_list)) == 1 and not isinstance(pooled_data, list):
81 + # if all dimensions are the same
82 + d = dimensions_list[0]
83 + pooled_data = pooled_data[..., :d]
84 + else:
85 + pooled_data = [
86 + vecs if d is None else vecs[..., :d]
87 + for vecs, d in zip(pooled_data, dimensions_list)
88 + ]
89 +
90 + # for normalize
91 + flags = [p.normalize for p in pooling_params]
92 + if len(set(flags)) == 1:
93 + if flags[0]:
94 + pooled_data = self.activation(pooled_data)
95 + else:
96 + pooled_data = [
97 + self.activation(vecs) if f else vecs
98 + for vecs, f in zip(pooled_data, flags)
99 + ]
100 +
101 + # pooled_data shape: [batchsize, embedding_dimension]
102 + return pooled_data
103 +
104 +
105 +class ClassifierPoolerHead(SequencePoolerHead):
106 + def __init__(
107 + self,
108 + classifier: ClassifierFn | None = None,
109 + act_fn: PoolerActivation | str | None = None,
110 + ) -> None:
111 + super().__init__()
112 +
113 + vllm_config = get_current_vllm_config()
114 + model_config = vllm_config.model_config
115 +
116 + self.classifier = classifier
117 + self.logit_bias: float | None = model_config.pooler_config.logit_bias
118 + self.head_dtype = model_config.head_dtype
119 +
120 + self.act_fn = resolve_classifier_act_fn(
121 + model_config, static_num_labels=True, act_fn=act_fn
122 + )
123 +
124 + def get_supported_tasks(self) -> Set[PoolingTask]:
125 + return {"classify", "score"}
126 +
127 + def forward(
128 + self,
129 + pooled_data: SequencePoolingMethodOutput,
130 + pooling_metadata: PoolingMetadata,
131 + ) -> SequencePoolerHeadOutput:
132 + pooling_params = pooling_metadata.pooling_params
133 + assert len(pooled_data) == len(pooling_params)
134 +
135 + if isinstance(pooled_data, list):
136 + pooled_data = torch.stack(pooled_data)
137 + # pooled_data shape: [batchsize, hidden_size]
138 +
139 + pooled_data = pooled_data.to(self.head_dtype)
140 +
141 + if self.classifier is not None:
142 + pooled_data = self.classifier(pooled_data)
143 + # pooled_data shape: [batchsize, num_labels]
144 +
145 + if self.logit_bias is not None:
146 + pooled_data -= self.logit_bias
147 +
148 + flags = [p.use_activation for p in pooling_params]
149 + if len(set(flags)) == 1:
150 + scores = self.act_fn(pooled_data) if flags[0] else pooled_data
151 + else:
152 + scores = [
153 + self.act_fn(vecs) if f else vecs for vecs, f in zip(pooled_data, flags)
154 + ]
155 +
156 + # scores shape: [batchsize, num_labels]
157 + return scores
vllm/model_executor/layers/pooler/seqwise/methods.py new
+93
@@ -0,0 +1,93 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Set
5 +from typing import TypeAlias
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config.pooler import PoolingTypeStr
11 +from vllm.model_executor.layers.pooler import PoolingParamsUpdate
12 +from vllm.tasks import PoolingTask
13 +from vllm.v1.pool.metadata import PoolingMetadata
14 +
15 +SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
16 +
17 +
18 +class SequencePoolingMethod(nn.Module, ABC):
19 + def get_supported_tasks(self) -> Set[PoolingTask]:
20 + return {"token_embed", "token_classify", "embed", "classify", "score"}
21 +
22 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
23 + return PoolingParamsUpdate()
24 +
25 + @abstractmethod
26 + def forward(
27 + self,
28 + hidden_states: torch.Tensor,
29 + pooling_metadata: PoolingMetadata,
30 + ) -> SequencePoolingMethodOutput:
31 + raise NotImplementedError
32 +
33 +
34 +class CLSPool(SequencePoolingMethod):
35 + def forward(
36 + self,
37 + hidden_states: torch.Tensor,
38 + pooling_metadata: PoolingMetadata,
39 + ) -> SequencePoolingMethodOutput:
40 + pooling_cursor = pooling_metadata.get_pooling_cursor()
41 + assert not pooling_cursor.is_partial_prefill(), (
42 + "partial prefill not supported with CLS pooling"
43 + )
44 +
45 + return hidden_states[pooling_cursor.first_token_indices_gpu]
46 +
47 +
48 +class LastPool(SequencePoolingMethod):
49 + def forward(
50 + self,
51 + hidden_states: torch.Tensor,
52 + pooling_metadata: PoolingMetadata,
53 + ) -> SequencePoolingMethodOutput:
54 + pooling_cursor = pooling_metadata.get_pooling_cursor()
55 + return hidden_states[pooling_cursor.last_token_indices_gpu]
56 +
57 +
58 +class MeanPool(SequencePoolingMethod):
59 + def forward(
60 + self,
61 + hidden_states: torch.Tensor,
62 + pooling_metadata: PoolingMetadata,
63 + ) -> SequencePoolingMethodOutput:
64 + pooling_cursor = pooling_metadata.get_pooling_cursor()
65 + assert not pooling_cursor.is_partial_prefill(), (
66 + "partial prefill not supported with MEAN pooling"
67 + )
68 +
69 + prompt_lens = pooling_cursor.prompt_lens_cpu.to(
70 + hidden_states.device, non_blocking=True
71 + )
72 +
73 + # Use float32 for torch.cumsum in MeanPool,
74 + # otherwise precision will be lost significantly.
75 + cumsum = torch.cumsum(hidden_states, dim=0, dtype=torch.float32)
76 +
77 + start_indices = pooling_cursor.first_token_indices_gpu
78 + end_indices = pooling_cursor.last_token_indices_gpu
79 +
80 + return (
81 + cumsum[end_indices] - cumsum[start_indices] + hidden_states[start_indices]
82 + ) / prompt_lens.unsqueeze(1)
83 +
84 +
85 +def get_seq_pooling_method(pooling_type: PoolingTypeStr | str):
86 + if pooling_type == "LAST":
87 + return LastPool()
88 + if pooling_type == "CLS":
89 + return CLSPool()
90 + if pooling_type == "MEAN":
91 + return MeanPool()
92 +
93 + raise NotImplementedError(f"Unknown sequence pooling type: {pooling_type!r}")
vllm/model_executor/layers/pooler/seqwise/poolers.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable, Set
4 +from typing import TypeAlias
5 +
6 +import torch
7 +
8 +from vllm.config import PoolerConfig
9 +from vllm.model_executor.layers.pooler import ClassifierFn, PoolingParamsUpdate
10 +from vllm.model_executor.layers.pooler.abstract import Pooler
11 +from vllm.model_executor.layers.pooler.activations import PoolerActivation
12 +from vllm.tasks import POOLING_TASKS, PoolingTask
13 +from vllm.v1.pool.metadata import PoolingMetadata
14 +
15 +from .heads import (
16 + ClassifierPoolerHead,
17 + EmbeddingPoolerHead,
18 + SequencePoolerHead,
19 + SequencePoolerHeadOutput,
20 +)
21 +from .methods import (
22 + SequencePoolingMethod,
23 + SequencePoolingMethodOutput,
24 + get_seq_pooling_method,
25 +)
26 +
27 +SequencePoolingFn: TypeAlias = Callable[
28 + [torch.Tensor, PoolingMetadata],
29 + SequencePoolingMethodOutput,
30 +]
31 +SequencePoolingHeadFn: TypeAlias = Callable[
32 + [SequencePoolingMethodOutput, PoolingMetadata],
33 + SequencePoolerHeadOutput,
34 +]
35 +
36 +SequencePoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
37 +
38 +
39 +class SequencePooler(Pooler):
40 + """
41 + A layer that pools specific information from hidden states.
42 +
43 + This layer does the following:
44 + 1. Extracts specific tokens or aggregates data based on pooling method.
45 + 2. Postprocesses the output based on pooling head.
46 + 3. Returns structured results as `PoolerOutput`.
47 + """
48 +
49 + def __init__(
50 + self,
51 + pooling: SequencePoolingMethod | SequencePoolingFn,
52 + head: SequencePoolerHead | SequencePoolingHeadFn,
53 + ) -> None:
54 + super().__init__()
55 +
56 + self.pooling = pooling
57 + self.head = head
58 +
59 + def get_supported_tasks(self) -> Set[PoolingTask]:
60 + tasks = set(POOLING_TASKS)
61 +
62 + if isinstance(self.pooling, SequencePoolingMethod):
63 + tasks &= self.pooling.get_supported_tasks()
64 + if isinstance(self.head, SequencePoolerHead):
65 + tasks &= self.head.get_supported_tasks()
66 +
67 + return tasks
68 +
69 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
70 + updates = PoolingParamsUpdate()
71 +
72 + if isinstance(self.pooling, SequencePoolingMethod):
73 + updates |= self.pooling.get_pooling_updates(task)
74 +
75 + return updates
76 +
77 + def forward(
78 + self,
79 + hidden_states: torch.Tensor,
80 + pooling_metadata: PoolingMetadata,
81 + ) -> SequencePoolerOutput:
82 + pooled_data = self.pooling(hidden_states, pooling_metadata)
83 + pooled_data = self.head(pooled_data, pooling_metadata)
84 + return pooled_data
85 +
86 +
87 +def pooler_for_embed(pooler_config: PoolerConfig):
88 + pooling = get_seq_pooling_method(pooler_config.get_pooling_type())
89 + head = EmbeddingPoolerHead()
90 +
91 + return SequencePooler(pooling=pooling, head=head)
92 +
93 +
94 +def pooler_for_classify(
95 + pooler_config: PoolerConfig,
96 + *,
97 + pooling: SequencePoolingMethod | SequencePoolingFn | None = None,
98 + classifier: ClassifierFn | None = None,
99 + act_fn: PoolerActivation | str | None = None,
100 +):
101 + if pooling is None:
102 + pooling = get_seq_pooling_method(pooler_config.get_pooling_type())
103 +
104 + head = ClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
105 +
106 + return SequencePooler(pooling=pooling, head=head)
vllm/model_executor/layers/pooler/special.py new
+128
@@ -0,0 +1,128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Mapping, Set
4 +from itertools import groupby
5 +
6 +import torch
7 +
8 +from vllm.config import PoolerConfig
9 +from vllm.model_executor.layers.pooler import PoolingParamsUpdate
10 +from vllm.tasks import PoolingTask
11 +from vllm.v1.pool.metadata import PoolingMetadata
12 +
13 +from .abstract import Pooler, PoolerOutput
14 +from .common import ClassifierFn
15 +from .seqwise import (
16 + SequencePoolingFn,
17 + SequencePoolingMethod,
18 + pooler_for_classify,
19 + pooler_for_embed,
20 +)
21 +from .tokwise import AllPool, pooler_for_token_classify, pooler_for_token_embed
22 +
23 +
24 +class DispatchPooler(Pooler):
25 + """Dispatches calls to a sub-pooler based on the pooling task."""
26 +
27 + @classmethod
28 + def for_embedding(cls, pooler_config: PoolerConfig):
29 + return cls(
30 + {
31 + "token_embed": pooler_for_token_embed(pooler_config),
32 + "embed": pooler_for_embed(pooler_config),
33 + },
34 + )
35 +
36 + @classmethod
37 + def for_seq_cls(
38 + cls,
39 + pooler_config: PoolerConfig,
40 + *,
41 + pooling: SequencePoolingMethod | SequencePoolingFn | None = None,
42 + classifier: ClassifierFn | None = None,
43 + ):
44 + return cls(
45 + {
46 + "token_classify": pooler_for_token_classify(
47 + pooler_config,
48 + pooling=AllPool(),
49 + classifier=classifier,
50 + ),
51 + "classify": pooler_for_classify(
52 + pooler_config,
53 + pooling=pooling,
54 + classifier=classifier,
55 + act_fn="classify",
56 + ),
57 + "score": pooler_for_classify(
58 + pooler_config,
59 + pooling=pooling,
60 + classifier=classifier,
61 + act_fn="score",
62 + ),
63 + }
64 + )
65 +
66 + def __init__(self, poolers_by_task: Mapping[PoolingTask, Pooler]) -> None:
67 + super().__init__()
68 +
69 + for task, pooler in poolers_by_task.items():
70 + if task not in pooler.get_supported_tasks():
71 + raise ValueError(
72 + f"{pooler=} does not support {task=}. "
73 + f"Supported tasks: {pooler.get_supported_tasks()}"
74 + )
75 +
76 + self.poolers_by_task = poolers_by_task
77 +
78 + def get_supported_tasks(self) -> Set[PoolingTask]:
79 + return set(self.poolers_by_task)
80 +
81 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
82 + return self.poolers_by_task[task].get_pooling_updates(task)
83 +
84 + def forward(
85 + self,
86 + hidden_states: torch.Tensor,
87 + pooling_metadata: PoolingMetadata,
88 + ) -> PoolerOutput:
89 + poolers_by_task = self.poolers_by_task
90 +
91 + outputs = list[torch.Tensor | None]()
92 + offset = 0
93 + for task, group in groupby(pooling_metadata.tasks):
94 + if not (pooler := poolers_by_task.get(task)):
95 + raise ValueError(
96 + f"Unsupported task: {task!r} "
97 + f"Supported tasks: {self.get_supported_tasks()}"
98 + )
99 +
100 + num_items = len(list(group))
101 + group_output: PoolerOutput = pooler(
102 + hidden_states,
103 + pooling_metadata[offset : offset + num_items],
104 + )
105 +
106 + outputs.extend(group_output)
107 + offset += num_items
108 +
109 + return outputs
110 +
111 + def extra_repr(self) -> str:
112 + s = f"supported_task={self.get_supported_tasks()}"
113 + return s
114 +
115 +
116 +class IdentityPooler(Pooler):
117 + def get_supported_tasks(self) -> Set[PoolingTask]:
118 + return {"plugin", "score"}
119 +
120 + def forward(
121 + self,
122 + hidden_states: torch.Tensor,
123 + pooling_metadata: PoolingMetadata,
124 + ) -> PoolerOutput:
125 + return hidden_states
126 +
127 +
128 +__all__ = ["DispatchPooler", "IdentityPooler"]
vllm/model_executor/layers/pooler/tokwise/__init__.py new
+39
@@ -0,0 +1,39 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Poolers that produce an output for each token in the sequence."""
4 +
5 +from .heads import (
6 + TokenClassifierPoolerHead,
7 + TokenEmbeddingPoolerHead,
8 + TokenPoolerHead,
9 + TokenPoolerHeadOutputItem,
10 +)
11 +from .methods import (
12 + AllPool,
13 + StepPool,
14 + TokenPoolingMethod,
15 + TokenPoolingMethodOutputItem,
16 + get_tok_pooling_method,
17 +)
18 +from .poolers import (
19 + TokenPooler,
20 + TokenPoolerOutput,
21 + pooler_for_token_classify,
22 + pooler_for_token_embed,
23 +)
24 +
25 +__all__ = [
26 + "TokenPoolerHead",
27 + "TokenPoolerHeadOutputItem",
28 + "TokenClassifierPoolerHead",
29 + "TokenEmbeddingPoolerHead",
30 + "TokenPoolingMethod",
31 + "TokenPoolingMethodOutputItem",
32 + "AllPool",
33 + "StepPool",
34 + "get_tok_pooling_method",
35 + "TokenPooler",
36 + "TokenPoolerOutput",
37 + "pooler_for_token_classify",
38 + "pooler_for_token_embed",
39 +]
vllm/model_executor/layers/pooler/tokwise/heads.py new
+142
@@ -0,0 +1,142 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Set
5 +from typing import TypeAlias
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import get_current_vllm_config
11 +from vllm.model_executor.layers.pooler import ClassifierFn
12 +from vllm.model_executor.layers.pooler.activations import (
13 + PoolerActivation,
14 + PoolerNormalize,
15 + resolve_classifier_act_fn,
16 +)
17 +from vllm.model_executor.models.adapters import _load_st_projector
18 +from vllm.pooling_params import PoolingParams
19 +from vllm.tasks import PoolingTask
20 +from vllm.v1.pool.metadata import PoolingMetadata
21 +
22 +from .methods import TokenPoolingMethodOutputItem
23 +
24 +TokenPoolerHeadOutputItem: TypeAlias = torch.Tensor | None
25 +
26 +
27 +class TokenPoolerHead(nn.Module, ABC):
28 + @abstractmethod
29 + def get_supported_tasks(self) -> Set[PoolingTask]:
30 + raise NotImplementedError
31 +
32 + @abstractmethod
33 + def forward_chunk(
34 + self,
35 + pooled_data: TokenPoolingMethodOutputItem,
36 + pooling_param: PoolingParams,
37 + ) -> TokenPoolerHeadOutputItem:
38 + raise NotImplementedError
39 +
40 + def forward(
41 + self,
42 + pooled_data: list[TokenPoolingMethodOutputItem],
43 + pooling_metadata: PoolingMetadata,
44 + ) -> list[TokenPoolerHeadOutputItem]:
45 + pooling_params = pooling_metadata.pooling_params
46 + assert len(pooled_data) == len(pooling_params)
47 +
48 + return [self.forward_chunk(d, p) for d, p in zip(pooled_data, pooling_params)]
49 +
50 +
51 +class TokenEmbeddingPoolerHead(TokenPoolerHead):
52 + def __init__(self) -> None:
53 + super().__init__()
54 +
55 + # Load ST projector if available
56 + vllm_config = get_current_vllm_config()
57 + model_config = vllm_config.model_config
58 +
59 + self.projector = _load_st_projector(model_config)
60 + self.head_dtype = model_config.head_dtype
61 +
62 + self.activation = PoolerNormalize()
63 +
64 + def get_supported_tasks(self) -> Set[PoolingTask]:
65 + return {"token_embed"}
66 +
67 + def forward_chunk(
68 + self,
69 + pooled_data: TokenPoolingMethodOutputItem,
70 + pooling_param: PoolingParams,
71 + ) -> TokenPoolerHeadOutputItem:
72 + # for unfinished chunked prefill
73 + if pooled_data is None:
74 + return None
75 +
76 + pooled_data = pooled_data.to(self.head_dtype)
77 + # pooled_data shape: [n_tokens, hidden_dimension]
78 +
79 + # Apply ST projector
80 + if self.projector is not None:
81 + pooled_data = self.projector(pooled_data)
82 + # pooled_data shape: [n_tokens, embedding_dimension]
83 +
84 + # for matryoshka representation
85 + pooled_data = pooled_data[..., : pooling_param.dimensions]
86 +
87 + # for normalize
88 + if pooling_param.normalize:
89 + pooled_data = self.activation(pooled_data)
90 +
91 + # pooled_data shape: [n_tokens, embedding_dimension]
92 + return pooled_data
93 +
94 +
95 +class TokenClassifierPoolerHead(TokenPoolerHead):
96 + def __init__(
97 + self,
98 + classifier: ClassifierFn | None = None,
99 + act_fn: PoolerActivation | str | None = None,
100 + ) -> None:
101 + super().__init__()
102 +
103 + vllm_config = get_current_vllm_config()
104 + model_config = vllm_config.model_config
105 +
106 + self.classifier = classifier
107 + self.logit_bias: float | None = model_config.pooler_config.logit_bias
108 + self.head_dtype = model_config.head_dtype
109 +
110 + self.act_fn = resolve_classifier_act_fn(
111 + model_config, static_num_labels=False, act_fn=act_fn
112 + )
113 +
114 + def get_supported_tasks(self) -> Set[PoolingTask]:
115 + return {"token_classify"}
116 +
117 + def forward_chunk(
118 + self,
119 + pooled_data: TokenPoolingMethodOutputItem,
120 + pooling_param: PoolingParams,
121 + ) -> TokenPoolerHeadOutputItem:
122 + # for unfinished chunked prefill
123 + if pooled_data is None:
124 + return None
125 +
126 + pooled_data = pooled_data.to(self.head_dtype)
127 + # hidden_states shape: [n_token, hidden_size]
128 +
129 + if self.classifier is not None:
130 + scores = self.classifier(pooled_data)
131 + else:
132 + scores = pooled_data
133 + # scores shape: [n_token, num_labels]
134 +
135 + if self.logit_bias is not None:
136 + scores -= self.logit_bias
137 +
138 + if pooling_param.use_activation:
139 + scores = self.act_fn(scores)
140 +
141 + # scores shape: [n_token, num_labels]
142 + return scores
vllm/model_executor/layers/pooler/tokwise/methods.py new
+124
@@ -0,0 +1,124 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Set
5 +from typing import TypeAlias
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import get_current_vllm_config
11 +from vllm.config.pooler import PoolingTypeStr
12 +from vllm.model_executor.layers.pooler import PoolingParamsUpdate
13 +from vllm.tasks import PoolingTask
14 +from vllm.v1.pool.metadata import PoolingMetadata
15 +
16 +TokenPoolingMethodOutputItem: TypeAlias = torch.Tensor | None
17 +
18 +
19 +class TokenPoolingMethod(nn.Module, ABC):
20 + def get_supported_tasks(self) -> Set[PoolingTask]:
21 + return {"token_embed", "token_classify"}
22 +
23 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
24 + return PoolingParamsUpdate()
25 +
26 + @abstractmethod
27 + def forward(
28 + self,
29 + hidden_states: torch.Tensor,
30 + pooling_metadata: PoolingMetadata,
31 + ) -> list[TokenPoolingMethodOutputItem]:
32 + raise NotImplementedError
33 +
34 +
35 +class AllPool(TokenPoolingMethod):
36 + def __init__(self):
37 + super().__init__()
38 +
39 + vllm_config = get_current_vllm_config()
40 + scheduler_config = vllm_config.scheduler_config
41 +
42 + self.enable_chunked_prefill = scheduler_config.enable_chunked_prefill
43 +
44 + def forward(
45 + self,
46 + hidden_states: torch.Tensor,
47 + pooling_metadata: PoolingMetadata,
48 + ) -> list[TokenPoolingMethodOutputItem]:
49 + pooling_cursor = pooling_metadata.get_pooling_cursor()
50 + hidden_states_all = hidden_states.split(
51 + pooling_cursor.num_scheduled_tokens_cpu.tolist()
52 + )
53 + hidden_states_lst = [hidden_states_all[i] for i in pooling_cursor.index]
54 +
55 + if not self.enable_chunked_prefill:
56 + return hidden_states_lst
57 +
58 + pooling_states = pooling_metadata.pooling_states
59 +
60 + # If chunked_prefill is enabled
61 + # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache
62 + for p, hs_chunk in zip(pooling_states, hidden_states_lst):
63 + p.hidden_states_cache.append(hs_chunk)
64 +
65 + # 2. Once prefill is finished, send hidden_states_cache to PoolerHead
66 + output_list = list[TokenPoolingMethodOutputItem]()
67 + for p, finished in zip(pooling_states, pooling_cursor.is_finished()):
68 + if finished:
69 + hidden_states_cache = p.hidden_states_cache
70 + if len(hidden_states_cache) == 1:
71 + output_list.append(hidden_states_cache[0])
72 + else:
73 + output_list.append(torch.concat(hidden_states_cache, dim=0))
74 + p.clean()
75 + else:
76 + output_list.append(None)
77 +
78 + return output_list
79 +
80 +
81 +class StepPool(AllPool):
82 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
83 + return PoolingParamsUpdate(requires_token_ids=True)
84 +
85 + def forward(
86 + self,
87 + hidden_states: torch.Tensor,
88 + pooling_metadata: PoolingMetadata,
89 + ) -> list[TokenPoolingMethodOutputItem]:
90 + pooled_data_lst = super().forward(hidden_states, pooling_metadata)
91 + prompt_token_ids = pooling_metadata.get_prompt_token_ids()
92 + pooling_params = pooling_metadata.pooling_params
93 +
94 + pooled_data = list[torch.Tensor | None]()
95 + for data, token_id, pooling_param in zip(
96 + pooled_data_lst, prompt_token_ids, pooling_params
97 + ):
98 + # for unfinished chunked prefill
99 + if data is None:
100 + pass
101 + else:
102 + step_tag_id = pooling_param.step_tag_id
103 + returned_token_ids = pooling_param.returned_token_ids
104 +
105 + if returned_token_ids is not None and len(returned_token_ids) > 0:
106 + data = data[:, returned_token_ids]
107 +
108 + if step_tag_id is not None:
109 + data = data[token_id == step_tag_id]
110 +
111 + pooled_data.append(data)
112 +
113 + return pooled_data
114 +
115 +
116 +def get_tok_pooling_method(pooling_type: PoolingTypeStr | str):
117 + if pooling_type == "ALL":
118 + return AllPool()
119 + if pooling_type == "STEP":
120 + return StepPool()
121 +
122 + # TODO: Separate seq and tok pooling types so we don't need this fallback
123 + return AllPool()
124 + raise NotImplementedError(f"Unknown tokenwise pooling type: {pooling_type!r}")
vllm/model_executor/layers/pooler/tokwise/poolers.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable, Set
4 +from typing import TypeAlias
5 +
6 +import torch
7 +
8 +from vllm.config import PoolerConfig
9 +from vllm.model_executor.layers.pooler import ClassifierFn, PoolingParamsUpdate
10 +from vllm.model_executor.layers.pooler.abstract import Pooler
11 +from vllm.model_executor.layers.pooler.activations import PoolerActivation
12 +from vllm.tasks import POOLING_TASKS, PoolingTask
13 +from vllm.v1.pool.metadata import PoolingMetadata
14 +
15 +from .heads import (
16 + TokenClassifierPoolerHead,
17 + TokenEmbeddingPoolerHead,
18 + TokenPoolerHead,
19 + TokenPoolerHeadOutputItem,
20 +)
21 +from .methods import (
22 + TokenPoolingMethod,
23 + TokenPoolingMethodOutputItem,
24 + get_tok_pooling_method,
25 +)
26 +
27 +TokenPoolingFn: TypeAlias = Callable[
28 + [torch.Tensor, PoolingMetadata],
29 + list[TokenPoolingMethodOutputItem],
30 +]
31 +TokenPoolingHeadFn: TypeAlias = Callable[
32 + [list[TokenPoolingMethodOutputItem], PoolingMetadata],
33 + list[TokenPoolerHeadOutputItem],
34 +]
35 +
36 +TokenPoolerOutput: TypeAlias = list[torch.Tensor | None]
37 +
38 +
39 +class TokenPooler(Pooler):
40 + """
41 + A layer that pools specific information from hidden states.
42 +
43 + This layer does the following:
44 + 1. Extracts specific tokens or aggregates data based on pooling method.
45 + 2. Postprocesses the output based on pooling head.
46 + 3. Returns structured results as `PoolerOutput`.
47 + """
48 +
49 + def __init__(
50 + self,
51 + pooling: TokenPoolingMethod | TokenPoolingFn,
52 + head: TokenPoolerHead | TokenPoolingHeadFn,
53 + ) -> None:
54 + super().__init__()
55 +
56 + self.pooling = pooling
57 + self.head = head
58 +
59 + def get_supported_tasks(self) -> Set[PoolingTask]:
60 + tasks = set(POOLING_TASKS)
61 +
62 + if isinstance(self.pooling, TokenPoolingMethod):
63 + tasks &= self.pooling.get_supported_tasks()
64 + if isinstance(self.head, TokenPoolerHead):
65 + tasks &= self.head.get_supported_tasks()
66 +
67 + return tasks
68 +
69 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
70 + updates = PoolingParamsUpdate()
71 +
72 + if isinstance(self.pooling, TokenPoolingMethod):
73 + updates |= self.pooling.get_pooling_updates(task)
74 +
75 + return updates
76 +
77 + def forward(
78 + self,
79 + hidden_states: torch.Tensor,
80 + pooling_metadata: PoolingMetadata,
81 + ) -> TokenPoolerOutput:
82 + pooled_data = self.pooling(hidden_states, pooling_metadata)
83 + pooled_data = self.head(pooled_data, pooling_metadata)
84 + return pooled_data
85 +
86 +
87 +def pooler_for_token_embed(pooler_config: PoolerConfig):
88 + pooling = get_tok_pooling_method(pooler_config.get_pooling_type())
89 + head = TokenEmbeddingPoolerHead()
90 +
91 + return TokenPooler(pooling=pooling, head=head)
92 +
93 +
94 +def pooler_for_token_classify(
95 + pooler_config: PoolerConfig,
96 + *,
97 + pooling: TokenPoolingMethod | TokenPoolingFn | None = None,
98 + classifier: ClassifierFn | None = None,
99 + act_fn: PoolerActivation | str | None = None,
100 +):
101 + if pooling is None:
102 + pooling = get_tok_pooling_method(pooler_config.get_pooling_type())
103 +
104 + head = TokenClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
105 +
106 + return TokenPooler(pooling=pooling, head=head)
vllm/model_executor/models/adapters.py
+5 -23
@@ -252,19 +252,14 @@ def as_embedding_model(cls: _T) -> _T:
252 return cls
253
254 # Lazy import
255 - from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
255 + from vllm.model_executor.layers.pooler import DispatchPooler
256
257 class ModelForEmbedding(_create_pooling_model_cls(cls)):
258 def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
259 pooler_config = vllm_config.model_config.pooler_config
260 assert pooler_config is not None
261
262 - self.pooler = DispatchPooler(
263 - {
264 - "token_embed": Pooler.for_token_embed(pooler_config),
265 - "embed": Pooler.for_embed(pooler_config),
266 - },
267 - )
262 + self.pooler = DispatchPooler.for_embedding(pooler_config)
263
264 ModelForEmbedding.__name__ = _get_pooling_model_name(cls.__name__, "ForEmbedding")
265
@@ -289,10 +284,7 @@ def as_seq_cls_model(cls: _T) -> _T:
284
285 # Lazy import
286 from vllm.model_executor.layers.linear import ReplicatedLinear
292 - from vllm.model_executor.layers.pooler import (
293 - DispatchPooler,
294 - Pooler,
295 - )
287 + from vllm.model_executor.layers.pooler import DispatchPooler
288 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
289
290 from .utils import maybe_prefix
@@ -318,18 +310,8 @@ def as_seq_cls_model(cls: _T) -> _T:
310 pooler_config = vllm_config.model_config.pooler_config
311 assert pooler_config is not None
312
321 - self.pooler = DispatchPooler(
322 - {
323 - "token_classify": Pooler.for_token_classify(
324 - pooler_config, classifier=self.score
325 - ),
326 - "classify": Pooler.for_classify(
327 - pooler_config, classifier=self.score, act_fn="classify"
328 - ),
329 - "score": Pooler.for_classify(
330 - pooler_config, classifier=self.score, act_fn="score"
331 - ),
332 - }
313 + self.pooler = DispatchPooler.for_seq_cls(
314 + pooler_config, classifier=self.score
315 )
316
317 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
vllm/model_executor/models/bert.py
+28 -57
@@ -18,19 +18,25 @@ from vllm.model_executor.layers.linear import (
18 RowParallelLinear,
19 )
20 from vllm.model_executor.layers.pooler import (
21 - ClassifierPooler,
21 DispatchPooler,
22 Pooler,
24 - PoolingMethod,
23 PoolingParamsUpdate,
26 - TokenPoolerHeadOutput,
27 - TokenPoolingMethodOutput,
24 +)
25 +from vllm.model_executor.layers.pooler.seqwise import (
26 + CLSPool,
27 + SequencePooler,
28 + SequencePoolerHeadOutput,
29 + SequencePoolerOutput,
30 + SequencePoolingMethodOutput,
31 +)
32 +from vllm.model_executor.layers.pooler.tokwise import (
33 + pooler_for_token_classify,
34 + pooler_for_token_embed,
35 )
36 from vllm.model_executor.layers.quantization import QuantizationConfig
37 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
38 from vllm.sequence import IntermediateTensors
39 from vllm.tasks import PoolingTask
33 -from vllm.v1.outputs import TokenPoolerOutput
40 from vllm.v1.pool.metadata import PoolingMetadata
41
42 from .interfaces import SupportsCrossEncoding, SupportsQuant
@@ -85,25 +91,21 @@ class BertEmbedding(nn.Module):
91 return embeddings
92
93
88 -class BertPooler(Pooler):
94 +class BertPooler(SequencePooler):
95 def __init__(self, config: BertConfig):
90 - super().__init__()
96 + super().__init__(
97 + pooling=CLSPool(),
98 + head=self.head,
99 + )
100
92 - self.pooling = PoolingMethod.from_pooling_type("CLS")
101 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
102 self.activation = nn.Tanh()
103
96 - def get_supported_tasks(self) -> Set[PoolingTask]:
97 - return self.pooling.get_supported_tasks()
98 -
99 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
100 - return self.pooling.get_pooling_updates(task)
101 -
104 def head(
105 self,
104 - pooled_data: TokenPoolingMethodOutput,
106 + pooled_data: SequencePoolingMethodOutput,
107 pooling_metadata: PoolingMetadata,
106 - ) -> TokenPoolerHeadOutput:
108 + ) -> SequencePoolerHeadOutput:
109 if isinstance(pooled_data, list):
110 pooled_data = torch.stack(pooled_data)
111
@@ -111,15 +113,6 @@ class BertPooler(Pooler):
113 pooled_data = self.activation(pooled_data)
114 return pooled_data
115
114 - def forward(
115 - self,
116 - hidden_states: torch.Tensor,
117 - pooling_metadata: PoolingMetadata,
118 - ) -> TokenPoolerOutput:
119 - pooled_data = self.pooling(hidden_states, pooling_metadata)
120 - pooled_data = self.head(pooled_data, pooling_metadata)
121 - return pooled_data
122 -
116
117 class BertEncoder(nn.Module):
118 def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
@@ -524,12 +517,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
517 )
518
519 def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
527 - return DispatchPooler(
528 - {
529 - "token_embed": Pooler.for_token_embed(pooler_config),
530 - "embed": Pooler.for_embed(pooler_config),
531 - }
532 - )
520 + return DispatchPooler.for_embedding(pooler_config)
521
522
523 # Here we encode the token type ids together with the input ids.
@@ -620,6 +608,7 @@ class SPLADESparsePooler(Pooler):
608 remove_cls_sep: bool = True,
609 ):
610 super().__init__()
611 +
612 assert pooling in ("max", "sum")
613 self.mlm_head = mlm_head
614 self.cls_token_id = cls_token_id
@@ -637,10 +626,8 @@ class SPLADESparsePooler(Pooler):
626 self,
627 hidden_states: torch.Tensor,
628 pooling_metadata: PoolingMetadata,
640 - ) -> torch.Tensor:
641 - assert isinstance(hidden_states, torch.Tensor) and hidden_states.dim() == 2
642 -
643 - lens_tensor: torch.Tensor = pooling_metadata.prompt_lens
629 + ) -> SequencePoolerOutput:
630 + lens_tensor = pooling_metadata.prompt_lens
631 lens: list[int] = lens_tensor.tolist()
632 B: int = len(lens)
633
@@ -729,7 +716,7 @@ class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):
716
717 return DispatchPooler(
718 {
732 - "token_embed": Pooler.for_token_embed(pooler_config),
719 + "token_embed": pooler_for_token_embed(pooler_config),
720 "embed": SPLADESparsePooler(
721 mlm_head=self.mlm_head,
722 cls_token_id=cls_id,
@@ -824,20 +811,10 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
811 pooler_config = vllm_config.model_config.pooler_config
812 assert pooler_config is not None
813
827 - self.pooler = DispatchPooler(
828 - {
829 - "token_classify": Pooler.for_token_classify(
830 - pooler_config, classifier=self.classifier
831 - ),
832 - "classify": ClassifierPooler(
833 - pooling=self.bert.pooler,
834 - classifier=self.classifier,
835 - act_fn="classify",
836 - ),
837 - "score": ClassifierPooler(
838 - pooling=self.bert.pooler, classifier=self.classifier, act_fn="score"
839 - ),
840 - }
814 + self.pooler = DispatchPooler.for_seq_cls(
815 + pooler_config,
816 + pooling=self.bert.pooler,
817 + classifier=self.classifier,
818 )
819
820 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -891,13 +868,7 @@ class BertForTokenClassification(nn.Module):
868 pooler_config = vllm_config.model_config.pooler_config
869 assert pooler_config is not None
870
894 - self.pooler = DispatchPooler(
895 - {
896 - "token_classify": Pooler.for_token_classify(
897 - pooler_config=pooler_config
898 - ),
899 - }
900 - )
871 + self.pooler = pooler_for_token_classify(pooler_config)
872
873 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
874 return self.bert.embed_input_ids(input_ids)
vllm/model_executor/models/bert_with_rope.py
+5 -15
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.linear import (
24 ReplicatedLinear,
25 RowParallelLinear,
26 )
27 +from vllm.model_executor.layers.pooler import DispatchPooler
28 from vllm.model_executor.layers.quantization import QuantizationConfig
29 from vllm.model_executor.layers.rotary_embedding import get_rope
30 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -37,7 +38,6 @@ from vllm.model_executor.utils import set_weight_attrs
38 from vllm.platforms import current_platform
39 from vllm.sequence import IntermediateTensors
40
40 -from ..layers.pooler import ClassifierPooler, DispatchPooler, Pooler
41 from .bert import BertPooler
42 from .interfaces import SupportsCrossEncoding, SupportsQuant
43 from .interfaces_base import default_pooling_type
@@ -693,20 +693,10 @@ class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):
693 pooler_config = vllm_config.model_config.pooler_config
694 assert pooler_config is not None
695
696 - self.pooler = DispatchPooler(
697 - {
698 - "token_classify": Pooler.for_token_classify(
699 - pooler_config, classifier=self.classifier
700 - ),
701 - "classify": ClassifierPooler(
702 - pooling=self.new.pooler,
703 - classifier=self.classifier,
704 - act_fn="classify",
705 - ),
706 - "score": ClassifierPooler(
707 - pooling=self.new.pooler, classifier=self.classifier, act_fn="score"
708 - ),
709 - }
696 + self.pooler = DispatchPooler.for_seq_cls(
697 + pooler_config,
698 + pooling=self.new.pooler,
699 + classifier=self.classifier,
700 )
701
702 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
vllm/model_executor/models/clip.py
+2 -7
@@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import (
26 QKVParallelLinear,
27 RowParallelLinear,
28 )
29 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
29 +from vllm.model_executor.layers.pooler import DispatchPooler
30 from vllm.model_executor.layers.quantization import QuantizationConfig
31 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
32 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -880,12 +880,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
880 assert pooler_config is not None
881 self.pooler_config = pooler_config
882
883 - self.pooler = DispatchPooler(
884 - {
885 - "token_embed": Pooler.for_token_embed(pooler_config),
886 - "embed": Pooler.for_embed(pooler_config),
887 - }
888 - )
883 + self.pooler = DispatchPooler.for_embedding(pooler_config)
884
885 # Assumes that self.forward is called after self.embed_input_ids
886 self._is_text_input = True
vllm/model_executor/models/gpt2.py
+2 -14
@@ -41,6 +41,7 @@ from vllm.model_executor.layers.linear import (
41 RowParallelLinear,
42 )
43 from vllm.model_executor.layers.logits_processor import LogitsProcessor
44 +from vllm.model_executor.layers.pooler import DispatchPooler
45 from vllm.model_executor.layers.quantization import QuantizationConfig
46 from vllm.model_executor.layers.vocab_parallel_embedding import (
47 ParallelLMHead,
@@ -49,7 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
50 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51 from vllm.sequence import IntermediateTensors
52
52 -from ..layers.pooler import DispatchPooler, Pooler
53 from .interfaces import SupportsCrossEncoding, SupportsPP
54 from .utils import (
55 AutoWeightsLoader,
@@ -351,19 +351,7 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):
351 pooler_config = vllm_config.model_config.pooler_config
352 assert pooler_config is not None
353
354 - self.pooler = DispatchPooler(
355 - {
356 - "token_classify": Pooler.for_token_classify(
357 - pooler_config, classifier=self.score
358 - ),
359 - "classify": Pooler.for_classify(
360 - pooler_config, classifier=self.score, act_fn="classify"
361 - ),
362 - "score": Pooler.for_classify(
363 - pooler_config, classifier=self.score, act_fn="score"
364 - ),
365 - }
366 - )
354 + self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
355
356 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
357 return self.transformer.embed_input_ids(input_ids)
vllm/model_executor/models/gritlm.py
+18 -29
@@ -9,17 +9,19 @@ from vllm.config import ModelConfig, VllmConfig
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.pooler import (
11 DispatchPooler,
12 - Pooler,
13 - PoolerNormalize,
14 - PoolingMethod,
12 PoolingParamsUpdate,
16 - TokenPoolerHeadOutput,
17 - TokenPoolingMethodOutput,
13 )
14 +from vllm.model_executor.layers.pooler.activations import PoolerNormalize
15 +from vllm.model_executor.layers.pooler.seqwise import (
16 + SequencePooler,
17 + SequencePoolerHeadOutput,
18 + SequencePoolingMethod,
19 + SequencePoolingMethodOutput,
20 +)
21 +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
22 from vllm.model_executor.models.llama import LlamaForCausalLM
23 from vllm.tasks import PoolingTask
24 from vllm.tokenizers import cached_tokenizer_from_config
22 -from vllm.v1.outputs import TokenPoolerOutput
25 from vllm.v1.pool.metadata import PoolingMetadata
26
27 from .interfaces_base import default_pooling_type
@@ -27,7 +29,7 @@ from .interfaces_base import default_pooling_type
29 logger = init_logger(__name__)
30
31
30 -class GritLMMeanPool(PoolingMethod):
32 +class GritLMMeanPool(SequencePoolingMethod):
33 """As `MeanPool`, but only includes non-instruction tokens."""
34
35 def __init__(self, model_config: ModelConfig):
@@ -151,7 +153,7 @@ class GritLMMeanPool(PoolingMethod):
153 self,
154 hidden_states: torch.Tensor,
155 pooling_metadata: PoolingMetadata,
154 - ) -> TokenPoolingMethodOutput:
156 + ) -> SequencePoolingMethodOutput:
157 prompt_lens = pooling_metadata.prompt_lens
158 instr_lens = torch.tensor(
159 [
@@ -174,35 +176,22 @@ class GritLMMeanPool(PoolingMethod):
176 return pooled_data
177
178
177 -class GritLMPooler(Pooler):
179 +class GritLMPooler(SequencePooler):
180 def __init__(self, model_config: ModelConfig):
179 - super().__init__()
181 + super().__init__(
182 + pooling=GritLMMeanPool(model_config),
183 + head=self.head,
184 + )
185
181 - self.pooling = GritLMMeanPool(model_config)
186 self.activation = PoolerNormalize()
187
184 - def get_supported_tasks(self) -> Set[PoolingTask]:
185 - return self.pooling.get_supported_tasks()
186 -
187 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
188 - return self.pooling.get_pooling_updates(task)
189 -
188 def head(
189 self,
192 - pooled_data: TokenPoolingMethodOutput,
190 + pooled_data: SequencePoolingMethodOutput,
191 pooling_metadata: PoolingMetadata,
194 - ) -> TokenPoolerHeadOutput:
192 + ) -> SequencePoolerHeadOutput:
193 return self.activation(pooled_data)
194
197 - def forward(
198 - self,
199 - hidden_states: torch.Tensor,
200 - pooling_metadata: PoolingMetadata,
201 - ) -> TokenPoolerOutput:
202 - pooled_data = self.pooling(hidden_states, pooling_metadata)
203 - pooled_data = self.head(pooled_data, pooling_metadata)
204 - return pooled_data
205 -
195
196 @default_pooling_type("MEAN")
197 class GritLM(LlamaForCausalLM):
@@ -245,7 +234,7 @@ class GritLM(LlamaForCausalLM):
234 if pooler_config is not None:
235 self.pooler = DispatchPooler(
236 {
248 - "token_embed": Pooler.for_token_embed(pooler_config),
237 + "token_embed": pooler_for_token_embed(pooler_config),
238 "embed": GritLMPooler(vllm_config.model_config),
239 }
240 )
vllm/model_executor/models/internlm2.py
+2 -4
@@ -28,7 +28,7 @@ from vllm.model_executor.layers.linear import (
28 RowParallelLinear,
29 )
30 from vllm.model_executor.layers.logits_processor import LogitsProcessor
31 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
31 +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
32 from vllm.model_executor.layers.quantization import QuantizationConfig
33 from vllm.model_executor.layers.rotary_embedding import get_rope
34 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -434,9 +434,7 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
434 pooler_config = vllm_config.model_config.pooler_config
435 assert pooler_config is not None
436
437 - self.pooler = DispatchPooler(
438 - {"token_classify": Pooler.for_token_classify(pooler_config)}
439 - )
437 + self.pooler = pooler_for_token_classify(pooler_config)
438
439 def forward(
440 self,
vllm/model_executor/models/jamba.py
+2 -14
@@ -27,7 +27,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
27 MambaStateDtypeCalculator,
28 MambaStateShapeCalculator,
29 )
30 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
30 +from vllm.model_executor.layers.pooler import DispatchPooler
31 from vllm.model_executor.layers.quantization import QuantizationConfig
32 from vllm.model_executor.layers.vocab_parallel_embedding import (
33 ParallelLMHead,
@@ -596,16 +596,4 @@ class JambaForSequenceClassification(JambaForCausalLM):
596 pooler_config = vllm_config.model_config.pooler_config
597 assert pooler_config is not None
598
599 - self.pooler = DispatchPooler(
600 - {
601 - "token_classify": Pooler.for_token_classify(
602 - pooler_config, classifier=self.score
603 - ),
604 - "classify": Pooler.for_classify(
605 - pooler_config, classifier=self.score, act_fn="classify"
606 - ),
607 - "score": Pooler.for_classify(
608 - pooler_config, classifier=self.score, act_fn="score"
609 - ),
610 - }
611 - )
599 + self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
vllm/model_executor/models/jina_vl.py
+2 -14
@@ -10,7 +10,7 @@ from vllm.config import ModelConfig, VllmConfig
10 from vllm.inputs import TokensPrompt
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
13 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
13 +from vllm.model_executor.layers.pooler import DispatchPooler
14 from vllm.multimodal import MULTIMODAL_REGISTRY
15 from vllm.sequence import IntermediateTensors
16
@@ -105,19 +105,7 @@ class JinaVLForSequenceClassification(
105 self.score = JinaVLScorer(
106 vllm_config.model_config, prefix=maybe_prefix(prefix, "score")
107 )
108 - self.pooler = DispatchPooler(
109 - {
110 - "token_classify": Pooler.for_token_classify(
111 - pooler_config, classifier=self.score
112 - ),
113 - "classify": Pooler.for_classify(
114 - pooler_config, classifier=self.score, act_fn="classify"
115 - ),
116 - "score": Pooler.for_classify(
117 - pooler_config, classifier=self.score, act_fn="score"
118 - ),
119 - }
120 - )
108 + self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
109
110 @classmethod
111 def get_placeholder_str(cls, modality: str, i: int) -> str | None:
vllm/model_executor/models/modernbert.py
+20 -51
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Iterable, Set
3 +from collections.abc import Iterable
4
5 import torch
6 from torch import nn
@@ -12,21 +12,18 @@ from vllm.compilation.decorators import support_torch_compile
12 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, RowParallelLinear
15 -from vllm.model_executor.layers.pooler import (
16 - ClassifierPooler,
17 - DispatchPooler,
18 - Pooler,
19 - PoolingMethod,
20 - PoolingParamsUpdate,
21 - TokenPoolerHeadOutput,
22 - TokenPoolingMethodOutput,
15 +from vllm.model_executor.layers.pooler import DispatchPooler
16 +from vllm.model_executor.layers.pooler.seqwise import (
17 + SequencePooler,
18 + SequencePoolerHeadOutput,
19 + SequencePoolingMethodOutput,
20 + get_seq_pooling_method,
21 )
22 +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
23 from vllm.model_executor.layers.rotary_embedding import get_rope
24 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
25 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
26 from vllm.sequence import IntermediateTensors
28 -from vllm.tasks import PoolingTask
29 -from vllm.v1.outputs import TokenPoolerOutput
27 from vllm.v1.pool.metadata import PoolingMetadata
28
29 from .interfaces import SupportsCrossEncoding
@@ -282,12 +279,13 @@ class ModernBertModel(nn.Module):
279 return norm_outputs
280
281
285 -class ModernBertPooler(Pooler):
282 +class ModernBertPooler(SequencePooler):
283 def __init__(self, config: ModernBertConfig):
287 - super().__init__()
284 + super().__init__(
285 + pooling=get_seq_pooling_method(config.classifier_pooling.upper()),
286 + head=self.head,
287 + )
288
289 - pooling_type = config.classifier_pooling.upper()
290 - self.pooling = PoolingMethod.from_pooling_type(pooling_type)
289 self.dense = nn.Linear(
290 config.hidden_size, config.hidden_size, config.classifier_bias
291 )
@@ -296,32 +294,17 @@ class ModernBertPooler(Pooler):
294 config.hidden_size, eps=config.norm_eps, bias=config.norm_bias
295 )
296
299 - def get_supported_tasks(self) -> Set[PoolingTask]:
300 - return self.pooling.get_supported_tasks()
301 -
302 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
303 - return self.pooling.get_pooling_updates(task)
304 -
297 def head(
298 self,
307 - pooled_data: TokenPoolingMethodOutput,
299 + pooled_data: SequencePoolingMethodOutput,
300 pooling_metadata: PoolingMetadata,
309 - ) -> TokenPoolerHeadOutput:
301 + ) -> SequencePoolerHeadOutput:
302 if isinstance(pooled_data, list):
303 pooled_data = torch.stack(pooled_data)
304
305 pooled_data = pooled_data.to(self.dense.weight.dtype)
306 return self.norm(self.act(self.dense(pooled_data)))
307
316 - def forward(
317 - self,
318 - hidden_states: torch.Tensor,
319 - pooling_metadata: PoolingMetadata,
320 - ) -> TokenPoolerOutput:
321 - pooled_data = self.pooling(hidden_states, pooling_metadata)
322 - pooled_data = self.head(pooled_data, pooling_metadata)
323 - return pooled_data
324 -
308
309 @default_pooling_type("CLS")
310 class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
@@ -344,18 +327,10 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
327 pooler_config = vllm_config.model_config.pooler_config
328 assert pooler_config is not None
329
347 - self.pooler = DispatchPooler(
348 - {
349 - "token_classify": Pooler.for_token_classify(
350 - pooler_config, classifier=self.classifier
351 - ),
352 - "classify": ClassifierPooler(
353 - pooling=self.pooling, classifier=self.classifier, act_fn="classify"
354 - ),
355 - "score": ClassifierPooler(
356 - pooling=self.pooling, classifier=self.classifier, act_fn="score"
357 - ),
358 - }
330 + self.pooler = DispatchPooler.for_seq_cls(
331 + pooler_config,
332 + pooling=self.pooling,
333 + classifier=self.classifier,
334 )
335
336 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -438,13 +413,7 @@ class ModernBertForTokenClassification(nn.Module):
413 pooler_config = vllm_config.model_config.pooler_config
414 assert pooler_config is not None
415
441 - self.pooler = DispatchPooler(
442 - {
443 - "token_classify": Pooler.for_token_classify(
444 - pooler_config=pooler_config
445 - ),
446 - }
447 - )
416 + self.pooler = pooler_for_token_classify(pooler_config)
417
418 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
419 return self.model.embed_input_ids(input_ids)
vllm/model_executor/models/qwen2_rm.py
+4 -7
@@ -14,7 +14,8 @@ from torch import nn
14
15 from vllm.config import VllmConfig
16 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
17 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
17 +from vllm.model_executor.layers.pooler import Pooler
18 +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
19 from vllm.sequence import IntermediateTensors
20
21 from .interfaces import SupportsLoRA, SupportsPP
@@ -104,9 +105,7 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
105 pooler_config = vllm_config.model_config.pooler_config
106 assert pooler_config is not None
107
107 - self.pooler = DispatchPooler(
108 - {"token_classify": Pooler.for_token_classify(pooler_config)}
109 - )
108 + self.pooler = pooler_for_token_classify(pooler_config)
109
110
111 @default_pooling_type("STEP")
@@ -118,6 +117,4 @@ class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
117 pooler_config = vllm_config.model_config.pooler_config
118 assert pooler_config is not None
119
121 - self.pooler = DispatchPooler(
122 - {"token_classify": Pooler.for_token_classify(pooler_config)}
123 - )
120 + self.pooler = pooler_for_token_classify(pooler_config)
vllm/model_executor/models/roberta.py
+6 -18
@@ -8,12 +8,8 @@ from torch import nn
8 from transformers import RobertaConfig
9
10 from vllm.config import ModelConfig, VllmConfig
11 -from vllm.model_executor.layers.pooler import (
12 - ClassifierPooler,
13 - CLSPool,
14 - DispatchPooler,
15 - Pooler,
16 -)
11 +from vllm.model_executor.layers.pooler import DispatchPooler
12 +from vllm.model_executor.layers.pooler.seqwise import CLSPool
13 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
14 from vllm.model_executor.models.bert import (
15 TOKEN_TYPE_SHIFT,
@@ -196,18 +192,10 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
192 pooler_config = vllm_config.model_config.pooler_config
193 assert pooler_config is not None
194
199 - self.pooler = DispatchPooler(
200 - {
201 - "token_classify": Pooler.for_token_classify(
202 - pooler_config=pooler_config, classifier=self.classifier
203 - ),
204 - "classify": ClassifierPooler(
205 - pooling=CLSPool(), classifier=self.classifier, act_fn="classify"
206 - ),
207 - "score": ClassifierPooler(
208 - pooling=CLSPool(), classifier=self.classifier, act_fn="score"
209 - ),
210 - }
195 + self.pooler = DispatchPooler.for_seq_cls(
196 + pooler_config,
197 + pooling=CLSPool(),
198 + classifier=self.classifier,
199 )
200
201 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
vllm/model_executor/models/siglip.py
+2 -7
@@ -27,7 +27,7 @@ from vllm.model_executor.layers.linear import (
27 QKVParallelLinear,
28 RowParallelLinear,
29 )
30 -from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
30 +from vllm.model_executor.layers.pooler import DispatchPooler
31 from vllm.model_executor.layers.quantization import QuantizationConfig
32 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
33 from vllm.model_executor.model_loader.weight_utils import (
@@ -1050,12 +1050,7 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1050 assert pooler_config is not None
1051 self.pooler_config = pooler_config
1052
1053 - self.pooler = DispatchPooler(
1054 - {
1055 - "token_embed": Pooler.for_token_embed(pooler_config),
1056 - "embed": Pooler.for_embed(pooler_config),
1057 - }
1058 - )
1053 + self.pooler = DispatchPooler.for_embedding(pooler_config)
1054
1055 self._is_text_input = True
1056
vllm/model_executor/models/terratorch.py
+2 -2
@@ -34,7 +34,7 @@ from transformers import BatchFeature
34 from vllm.config import VllmConfig
35 from vllm.config.multimodal import BaseDummyOptions
36 from vllm.logger import init_logger
37 -from vllm.model_executor.layers.pooler import DispatchPooler, DummyPooler
37 +from vllm.model_executor.layers.pooler import IdentityPooler
38 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39 from vllm.model_executor.models.utils import AutoWeightsLoader
40 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -248,7 +248,7 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
248 pooler_config = vllm_config.model_config.pooler_config
249 assert pooler_config is not None
250
251 - self.pooler = DispatchPooler({"plugin": DummyPooler()})
251 + self.pooler = IdentityPooler()
252
253 def embed_input_ids(
254 self,
vllm/model_executor/models/transformers/pooling.py
+7 -24
@@ -22,12 +22,8 @@ import torch
22 from transformers import AutoModelForSequenceClassification
23
24 from vllm.config.utils import getattr_iter
25 -from vllm.model_executor.layers.pooler import (
26 - ClassifierPooler,
27 - CLSPool,
28 - DispatchPooler,
29 - Pooler,
30 -)
25 +from vllm.model_executor.layers.pooler import DispatchPooler
26 +from vllm.model_executor.layers.pooler.seqwise import CLSPool
27 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
28 from vllm.model_executor.models.interfaces_base import VllmModelForPooling
29
@@ -47,12 +43,7 @@ class EmbeddingMixin(VllmModelForPooling):
43 pooler_config = vllm_config.model_config.pooler_config
44 assert pooler_config is not None
45
50 - self.pooler = DispatchPooler(
51 - {
52 - "token_embed": Pooler.for_token_embed(pooler_config),
53 - "embed": Pooler.for_embed(pooler_config),
54 - }
55 - )
46 + self.pooler = DispatchPooler.for_embedding(pooler_config)
47
48
49 class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):
@@ -104,16 +95,8 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling):
95
96 self.classifier.__class__ = ClassifierWithReshape
97
107 - self.pooler = DispatchPooler(
108 - {
109 - "token_classify": Pooler.for_token_classify(
110 - pooler_config, classifier=self.classifier
111 - ),
112 - "classify": ClassifierPooler(
113 - pooling=CLSPool(), classifier=self.classifier, act_fn="classify"
114 - ),
115 - "score": ClassifierPooler(
116 - pooling=CLSPool(), classifier=self.classifier, act_fn="score"
117 - ),
118 - }
98 + self.pooler = DispatchPooler.for_seq_cls(
99 + pooler_config,
100 + pooling=CLSPool(),
101 + classifier=self.classifier,
102 )
vllm/v1/outputs.py
+1 -3
@@ -91,9 +91,7 @@ class LogprobsTensors(NamedTuple):
91
92 # [num_reqs, <dynamic>]
93 # The shape of each element depends on the pooler used
94 -TokenPoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
95 -TokenwisePoolerOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]
96 -PoolerOutput: TypeAlias = TokenPoolerOutput | TokenwisePoolerOutput
94 +PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None]
95
96
97 @dataclass