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

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, 12:45 UTC

[Chore] Further cleanup pooler (#31951)

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

Cyrus Leung committed Jan 8, 2026 at 18:16 UTC d1b6fe007f1fd5afb5533a25dd188e1363a119ac

7 files changed +47 -62

tests/model_executor/test_model_load_with_params.py

+3 -8

@@ -5,12 +5,7 @@ import os

6 import pytest

8 -from vllm.model_executor.layers.pooler import (

9 - CLSPool,

10 - DispatchPooler,

11 - MeanPool,

12 - PoolingType,

13 -)

8 +from vllm.model_executor.layers.pooler import CLSPool, DispatchPooler, MeanPool

9 from vllm.model_executor.models.bert import BertEmbeddingModel

10 from vllm.model_executor.models.roberta import RobertaEmbeddingModel

11 from vllm.platforms import current_platform

@@ -50,7 +45,7 @@ def test_model_loading_with_params(vllm_runner, monkeypatch):

45 assert model_config.encoder_config["do_lower_case"]

47 # asserts on the pooling config files

53 - assert model_config.pooler_config.pooling_type == PoolingType.CLS.name

48 + assert model_config.pooler_config.pooling_type == "CLS"

49 assert model_config.pooler_config.normalize

51 # asserts on the tokenizer loaded

@@ -94,7 +89,7 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):

89 assert not model_config.encoder_config["do_lower_case"]

91 # asserts on the pooling config files

97 - assert model_config.pooler_config.pooling_type == PoolingType.MEAN.name

92 + assert model_config.pooler_config.pooling_type == "MEAN"

93 assert model_config.pooler_config.normalize

95 # asserts on the tokenizer loaded

tests/test_config.py

+1 -2

@@ -25,7 +25,6 @@ from vllm.config.vllm import (

25 OPTIMIZATION_LEVEL_TO_CONFIG,

26 OptimizationLevel,

27 )

28 -from vllm.model_executor.layers.pooler import PoolingType

28 from vllm.platforms import current_platform

@@ -162,7 +161,7 @@ def test_get_pooling_config():

161

162 assert model_config.pooler_config is not None

163 assert model_config.pooler_config.normalize

165 - assert model_config.pooler_config.pooling_type == PoolingType.MEAN.name

164 + assert model_config.pooler_config.pooling_type == "MEAN"

165

166

167 @pytest.mark.skipif(

vllm/config/pooler.py

+1 -2

@@ -21,8 +21,7 @@ class PoolerConfig:

22 pooling_type: PoolingTypeStr | None = None

23 """

24 - The pooling method of the pooling model. This should be a key in

25 - [`vllm.model_executor.layers.pooler.PoolingType`][].

24 + The pooling method of the pooling model.

25 """

27 ## for embeddings models

vllm/model_executor/layers/pooler.py

+38 -44

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

3 from abc import ABC, abstractmethod

4 from collections.abc import Callable, Mapping, Set

5 from dataclasses import dataclass

6 -from enum import IntEnum

6 from itertools import groupby

7 from typing import TypeAlias, TypeVar

@@ -12,13 +11,14 @@ import torch.nn as nn

11 import torch.nn.functional as F

12 from transformers import PretrainedConfig

15 -from vllm.config import ModelConfig, PoolerConfig, get_current_vllm_config

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, TokensPoolerOutput

21 +from vllm.v1.outputs import PoolerOutput, TokenPoolerOutput, TokenwisePoolerOutput

22 from vllm.v1.pool.metadata import PoolingMetadata

24 logger = init_logger(__name__)

@@ -31,27 +31,17 @@ ClassifierFn = Callable[[torch.Tensor], torch.Tensor]

33 TokenPoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]

34 -TokensPoolingMethodOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]

35 -TokensPoolingMethodOutputItem: TypeAlias = torch.Tensor | None

36 -PoolingMethodOutput: TypeAlias = TokenPoolingMethodOutput | TokensPoolingMethodOutput

34 +TokenwisePoolingMethodOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]

35 +TokenwisePoolingMethodOutputItem: TypeAlias = torch.Tensor | None

36 +PoolingMethodOutput: TypeAlias = TokenPoolingMethodOutput | TokenwisePoolingMethodOutput

38 TokenPoolerHeadOutput: TypeAlias = torch.Tensor | list[torch.Tensor]

39 -TokensPoolerHeadOutput: TypeAlias = torch.Tensor | None

40 -

41 -

42 -class PoolingType(IntEnum):

43 - """Enumeration for different types of pooling methods."""

44 -

45 - LAST = 0

46 - ALL = 1

47 - CLS = 2

48 - STEP = 3

49 - MEAN = 4

39 +TokenwisePoolerHeadOutput: TypeAlias = torch.Tensor | None

42 @dataclass(frozen=True)

43 class ResolvedPoolingConfig:

54 - pooling_type: PoolingType

44 + pooling_type: PoolingTypeStr

45 task: PoolingTask

47 @classmethod

@@ -61,7 +51,7 @@ class ResolvedPoolingConfig:

51 pooler_config: PoolerConfig,

52 ) -> "ResolvedPoolingConfig":

53 assert pooler_config.pooling_type is not None

64 - return cls(task=task, pooling_type=PoolingType[pooler_config.pooling_type])

54 + return cls(task=task, pooling_type=pooler_config.pooling_type)

57 @dataclass(frozen=True)

@@ -112,17 +102,22 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):

102

103 class PoolingMethod(nn.Module, ABC):

104 @staticmethod

115 - def from_pooling_type(pooling_type: PoolingType) -> "PoolingMethod":

116 - if pooling_type == PoolingType.LAST:

105 + def from_pooling_type(pooling_type: PoolingTypeStr) -> "PoolingMethod":

106 + if pooling_type == "LAST":

107 return LastPool()

118 - if pooling_type == PoolingType.ALL:

108 + if pooling_type == "ALL":

109 return AllPool()

120 - if pooling_type == PoolingType.CLS:

110 + if pooling_type == "CLS":

111 return CLSPool()

122 - if pooling_type == PoolingType.MEAN:

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

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

120 + raise NotImplementedError(f"Unsupported method: {pooling_type!r}")

121

122 @abstractmethod

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

@@ -186,13 +181,12 @@ class AllPool(PoolingMethod):

181 self,

182 hidden_states: torch.Tensor,

183 pooling_metadata: PoolingMetadata,

189 - ) -> TokensPoolingMethodOutput:

184 + ) -> TokenwisePoolingMethodOutput:

185 pooling_cursor = pooling_metadata.get_pooling_cursor()

191 - is_finished = pooling_cursor.is_finished()

192 - hidden_states_lst = list(

193 - hidden_states.split(pooling_cursor.num_scheduled_tokens_cpu.tolist())

186 + hidden_states_all = hidden_states.split(

187 + pooling_cursor.num_scheduled_tokens_cpu.tolist()

188 )

195 - hidden_states_lst = [hidden_states_lst[i] for i in pooling_cursor.index]

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

@@ -206,7 +200,7 @@ class AllPool(PoolingMethod):

200

201 # 2. Once prefill is finished, send hidden_states_cache to PoolerHead

202 output_list = list[torch.Tensor | None]()

209 - for p, finished in zip(pooling_states, is_finished):

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:

@@ -620,19 +614,19 @@ class ClassifierPooler(Pooler):

614 return scores

615

616

623 -class TokensPoolerHead(nn.Module, ABC):

617 +class TokenwisePoolerHead(nn.Module, ABC):

618 """Applicable to pooling strategies that output multiple tokens."""

619

620 @abstractmethod

621 def forward(

622 self,

629 - pooled_data: TokensPoolingMethodOutputItem,

623 + pooled_data: TokenwisePoolingMethodOutputItem,

624 pooling_param: PoolingParams,

631 - ) -> TokensPoolerHeadOutput:

625 + ) -> TokenwisePoolerHeadOutput:

626 raise NotImplementedError

627

628

635 -class TokenEmbeddingPoolerHead(TokensPoolerHead):

629 +class TokenEmbeddingPoolerHead(TokenwisePoolerHead):

630 def __init__(self) -> None:

631 super().__init__()

632

@@ -647,9 +641,9 @@ class TokenEmbeddingPoolerHead(TokensPoolerHead):

641

642 def forward(

643 self,

650 - pooled_data: TokensPoolingMethodOutputItem,

644 + pooled_data: TokenwisePoolingMethodOutputItem,

645 pooling_param: PoolingParams,

652 - ) -> TokensPoolerHeadOutput:

646 + ) -> TokenwisePoolerHeadOutput:

647 # for unfinished chunked prefill

648 if pooled_data is None:

649 return None

@@ -673,7 +667,7 @@ class TokenEmbeddingPoolerHead(TokensPoolerHead):

667 return pooled_data

668

669

676 -class TokenClassifierPoolerHead(TokensPoolerHead):

670 +class TokenClassifierPoolerHead(TokenwisePoolerHead):

671 def __init__(

672 self,

673 classifier: ClassifierFn | None,

@@ -695,9 +689,9 @@ class TokenClassifierPoolerHead(TokensPoolerHead):

689

690 def forward(

691 self,

698 - pooled_data: TokensPoolingMethodOutputItem,

692 + pooled_data: TokenwisePoolingMethodOutputItem,

693 pooling_param: PoolingParams,

700 - ) -> TokensPoolerHeadOutput:

694 + ) -> TokenwisePoolerHeadOutput:

695 # for unfinished chunked prefill

696 if pooled_data is None:

697 return None

@@ -722,7 +716,7 @@ class TokenClassifierPoolerHead(TokensPoolerHead):

716

717

718 class AllPooler(Pooler):

725 - def __init__(self, head: TokensPoolerHead) -> None:

719 + def __init__(self, head: TokenwisePoolerHead) -> None:

720 super().__init__()

721

722 self.pooling = AllPool()

@@ -735,7 +729,7 @@ class AllPooler(Pooler):

729 self,

730 hidden_states: torch.Tensor,

731 pooling_metadata: PoolingMetadata,

738 - ) -> TokensPoolerOutput:

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)

@@ -744,7 +738,7 @@ class AllPooler(Pooler):

738

739

740 class StepPooler(Pooler):

747 - def __init__(self, head: TokensPoolerHead) -> None:

741 + def __init__(self, head: TokenwisePoolerHead) -> None:

742 super().__init__()

743

744 self.pooling = AllPool()

@@ -790,7 +784,7 @@ class StepPooler(Pooler):

784 self,

785 hidden_states: torch.Tensor,

786 pooling_metadata: PoolingMetadata,

793 - ) -> TokensPoolerOutput:

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)

vllm/model_executor/models/bert.py

+1 -2

@@ -23,7 +23,6 @@ from vllm.model_executor.layers.pooler import (

23 Pooler,

24 PoolingMethod,

25 PoolingParamsUpdate,

26 - PoolingType,

26 TokenPoolerHeadOutput,

27 TokenPoolingMethodOutput,

28 )

@@ -90,7 +89,7 @@ class BertPooler(Pooler):

89 def __init__(self, config: BertConfig):

90 super().__init__()

93 - self.pooling = PoolingMethod.from_pooling_type(PoolingType.CLS)

92 + self.pooling = PoolingMethod.from_pooling_type("CLS")

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

94 self.activation = nn.Tanh()

vllm/model_executor/models/modernbert.py

+1 -2

@@ -18,7 +18,6 @@ from vllm.model_executor.layers.pooler import (

18 Pooler,

19 PoolingMethod,

20 PoolingParamsUpdate,

21 - PoolingType,

21 TokenPoolerHeadOutput,

22 TokenPoolingMethodOutput,

23 )

@@ -287,7 +286,7 @@ class ModernBertPooler(Pooler):

286 def __init__(self, config: ModernBertConfig):

287 super().__init__()

288

290 - pooling_type = PoolingType[config.classifier_pooling.upper()]

289 + pooling_type = config.classifier_pooling.upper()

290 self.pooling = PoolingMethod.from_pooling_type(pooling_type)

291 self.dense = nn.Linear(

292 config.hidden_size, config.hidden_size, config.classifier_bias

vllm/v1/outputs.py

+2 -2

@@ -92,8 +92,8 @@ class LogprobsTensors(NamedTuple):

92 # [num_reqs, <dynamic>]

93 # The shape of each element depends on the pooler used

94 TokenPoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor]

95 -TokensPoolerOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]

96 -PoolerOutput: TypeAlias = TokenPoolerOutput | TokensPoolerOutput

95 +TokenwisePoolerOutput: TypeAlias = list[torch.Tensor] | list[torch.Tensor | None]

96 +PoolerOutput: TypeAlias = TokenPoolerOutput | TokenwisePoolerOutput

99 @dataclass