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

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, 14:00 UTC

[Refactor] Clean up pooler modules (#31897)

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

Cyrus Leung committed Jan 8, 2026 at 00:07 UTC b7036c87a13bd94fabf9e46436d3c1e67688f729

7 files changed +167 -120

vllm/model_executor/layers/pooler.py

+102 -81

@@ -5,7 +5,7 @@ from collections.abc import Callable, Mapping, Set

5 from dataclasses import dataclass

6 from enum import IntEnum

7 from itertools import groupby

8 -from typing import TypeVar

8 +from typing import TypeAlias, TypeVar

10 import torch

11 import torch.nn as nn

@@ -18,8 +18,8 @@ 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

22 -from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata

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

22 +from vllm.v1.pool.metadata import PoolingMetadata

24 logger = init_logger(__name__)

@@ -30,6 +30,15 @@ PoolingFn = Callable[

30 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

37 +

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."""

@@ -123,31 +132,24 @@ class PoolingMethod(nn.Module, ABC):

132 return PoolingParamsUpdate()

133

134 @abstractmethod

126 - def forward_all(

127 - self,

128 - hidden_states: torch.Tensor,

129 - pooling_cursor: PoolingCursor,

130 - ) -> PoolerOutput:

131 - raise NotImplementedError

132 -

135 def forward(

136 self,

137 hidden_states: torch.Tensor,

138 pooling_metadata: PoolingMetadata,

137 - ) -> PoolerOutput:

138 - pooling_cursor = pooling_metadata.pooling_cursor

139 - return self.forward_all(hidden_states, pooling_cursor)

139 + ) -> PoolingMethodOutput:

140 + raise NotImplementedError

141

142

143 class CLSPool(PoolingMethod):

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

145 return {"token_embed", "token_classify", "embed", "classify", "score"}

146

146 - def forward_all(

147 + def forward(

148 self,

149 hidden_states: torch.Tensor,

149 - pooling_cursor: PoolingCursor,

150 - ) -> PoolerOutput:

150 + pooling_metadata: PoolingMetadata,

151 + ) -> TokenPoolingMethodOutput:

152 + pooling_cursor = pooling_metadata.get_pooling_cursor()

153 assert not pooling_cursor.is_partial_prefill(), (

154 "partial prefill not supported with CLS pooling"

155 )

@@ -159,11 +161,12 @@ class LastPool(PoolingMethod):

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

162 return {"token_embed", "token_classify", "embed", "classify", "score"}

163

162 - def forward_all(

164 + def forward(

165 self,

166 hidden_states: torch.Tensor,

165 - pooling_cursor: PoolingCursor,

166 - ) -> PoolerOutput:

167 + pooling_metadata: PoolingMetadata,

168 + ) -> TokenPoolingMethodOutput:

169 + pooling_cursor = pooling_metadata.get_pooling_cursor()

170 return hidden_states[pooling_cursor.last_token_indices_gpu]

171

172

@@ -179,19 +182,12 @@ class AllPool(PoolingMethod):

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

183 return {"token_embed", "token_classify"}

184

182 - def forward_all(

183 - self, hidden_states: torch.Tensor, pooling_cursor: PoolingCursor

184 - ) -> PoolerOutput:

185 - raise NotImplementedError(

186 - "forward_all is not implemented for AllPool. Use forward instead."

187 - )

188 -

185 def forward(

186 self,

187 hidden_states: torch.Tensor,

188 pooling_metadata: PoolingMetadata,

193 - ) -> PoolerOutput:

194 - pooling_cursor = pooling_metadata.pooling_cursor

189 + ) -> TokensPoolingMethodOutput:

190 + 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())

@@ -209,7 +205,7 @@ class AllPool(PoolingMethod):

205 p.hidden_states_cache.append(hs_chunk)

206

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

212 - output_list: PoolerOutput = []

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

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

210 if finished:

211 hidden_states_cache = p.hidden_states_cache

@@ -228,11 +224,12 @@ class MeanPool(PoolingMethod):

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

225 return {"token_embed", "token_classify", "embed", "classify", "score"}

226

231 - def forward_all(

227 + def forward(

228 self,

229 hidden_states: torch.Tensor,

234 - pooling_cursor: PoolingCursor,

235 - ) -> PoolerOutput:

230 + pooling_metadata: PoolingMetadata,

231 + ) -> TokenPoolingMethodOutput:

232 + pooling_cursor = pooling_metadata.get_pooling_cursor()

233 assert not pooling_cursor.is_partial_prefill(), (

234 "partial prefill not supported with MEAN pooling"

235 )

@@ -410,7 +407,7 @@ class Pooler(nn.Module, ABC):

407 @abstractmethod

408 def forward(

409 self,

413 - hidden_states: list[torch.Tensor] | torch.Tensor,

410 + hidden_states: torch.Tensor,

411 pooling_metadata: PoolingMetadata,

412 ) -> PoolerOutput:

413 raise NotImplementedError

@@ -422,41 +419,42 @@ class DummyPooler(Pooler):

419

420 def forward(

421 self,

425 - hidden_states: list[torch.Tensor] | torch.Tensor,

422 + hidden_states: torch.Tensor,

423 pooling_metadata: PoolingMetadata,

424 ) -> PoolerOutput:

425 return hidden_states

426

427

431 -class PoolerHead(nn.Module):

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

433 - super().__init__()

434 - self.activation = activation

428 +class TokenPoolerHead(nn.Module, ABC):

429 + """Applicable to pooling strategies that output one token."""

430

431 + @abstractmethod

432 def forward(

433 self,

438 - pooled_data: list[torch.Tensor] | torch.Tensor,

434 + pooled_data: TokenPoolingMethodOutput,

435 pooling_metadata: PoolingMetadata,

440 - ) -> PoolerOutput:

441 - return self.activation(pooled_data)

436 + ) -> TokenPoolerHeadOutput:

437 + raise NotImplementedError

438

439

444 -class EmbeddingPoolerHead(PoolerHead):

440 +class EmbeddingPoolerHead(TokenPoolerHead):

441 def __init__(self) -> None:

446 - super().__init__(activation=PoolerNormalize())

442 + super().__init__()

443

444 # Load ST projector if available

445 vllm_config = get_current_vllm_config()

450 - self.projector: nn.Module | None = (

446 + self.projector = (

447 _load_st_projector(vllm_config.model_config) if vllm_config else None

448 )

449 self.head_dtype = vllm_config.model_config.head_dtype

450

451 + self.activation = PoolerNormalize()

452 +

453 def forward(

454 self,

457 - pooled_data: list[torch.Tensor] | torch.Tensor,

455 + pooled_data: TokenPoolingMethodOutput,

456 pooling_metadata: PoolingMetadata,

459 - ) -> PoolerOutput:

457 + ) -> TokenPoolerHeadOutput:

458 if isinstance(pooled_data, list):

459 pooled_data = torch.stack(pooled_data)

460 # pooled_data shape: [batchsize, hidden_dimension]

@@ -509,7 +507,7 @@ class SimplePooler(Pooler):

507 3. Returns structured results as `PoolerOutput`.

508 """

509

512 - def __init__(self, pooling: PoolingMethod, head: PoolerHead) -> None:

510 + def __init__(self, pooling: PoolingMethod, head: TokenPoolerHead) -> None:

511 super().__init__()

512

513 self.pooling = pooling

@@ -523,9 +521,9 @@ class SimplePooler(Pooler):

521

522 def forward(

523 self,

526 - hidden_states: torch.Tensor | list[torch.Tensor],

524 + hidden_states: torch.Tensor,

525 pooling_metadata: PoolingMetadata,

528 - ) -> PoolerOutput:

526 + ) -> TokenPoolerHeadOutput:

527 pooled_data = self.pooling(hidden_states, pooling_metadata)

528 pooled_data = self.head(pooled_data, pooling_metadata)

529 return pooled_data

@@ -591,9 +589,9 @@ class ClassifierPooler(Pooler):

589

590 def forward(

591 self,

594 - hidden_states: torch.Tensor | list[torch.Tensor],

592 + hidden_states: torch.Tensor,

593 pooling_metadata: PoolingMetadata,

596 - ) -> PoolerOutput:

594 + ) -> TokenPoolerOutput:

595 pooled_data = self.pooling(hidden_states, pooling_metadata)

596 if isinstance(pooled_data, list):

597 pooled_data = torch.stack(pooled_data)

@@ -622,10 +620,36 @@ class ClassifierPooler(Pooler):

620 return scores

621

622

625 -class TokenEmbeddingPoolerHead(EmbeddingPoolerHead):

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

624 + """Applicable to pooling strategies that output multiple tokens."""

625 +

626 + @abstractmethod

627 def forward(

627 - self, pooled_data: torch.Tensor | None, pooling_param: PoolingParams

628 - ) -> PoolerOutput:

628 + self,

629 + pooled_data: TokensPoolingMethodOutputItem,

630 + pooling_param: PoolingParams,

631 + ) -> TokensPoolerHeadOutput:

632 + raise NotImplementedError

633 +

634 +

635 +class TokenEmbeddingPoolerHead(TokensPoolerHead):

636 + def __init__(self) -> None:

637 + super().__init__()

638 +

639 + # Load ST projector if available

640 + vllm_config = get_current_vllm_config()

641 + self.projector = (

642 + _load_st_projector(vllm_config.model_config) if vllm_config else None

643 + )

644 + self.head_dtype = vllm_config.model_config.head_dtype

645 +

646 + self.activation = PoolerNormalize()

647 +

648 + def forward(

649 + self,

650 + pooled_data: TokensPoolingMethodOutputItem,

651 + pooling_param: PoolingParams,

652 + ) -> TokensPoolerHeadOutput:

653 # for unfinished chunked prefill

654 if pooled_data is None:

655 return None

@@ -649,57 +673,56 @@ class TokenEmbeddingPoolerHead(EmbeddingPoolerHead):

673 return pooled_data

674

675

652 -class TokenClassifierPoolerHead(nn.Module):

676 +class TokenClassifierPoolerHead(TokensPoolerHead):

677 def __init__(

678 self,

679 classifier: ClassifierFn | None,

680 act_fn: PoolerActivation | str | None = None,

681 ) -> None:

682 super().__init__()

683 +

684 vllm_config = get_current_vllm_config()

685

686 self.classifier = classifier

662 - self.act_fn = ClassifierPooler.resolve_act_fn(

663 - vllm_config.model_config, static_num_labels=False, act_fn=act_fn

664 - )

687 self.logit_bias: float | None = (

688 vllm_config.model_config.pooler_config.logit_bias

689 )

690 self.head_dtype = vllm_config.model_config.head_dtype

691

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

671 - return {"token_classify"}

692 + self.activation = ClassifierPooler.resolve_act_fn(

693 + vllm_config.model_config, static_num_labels=False, act_fn=act_fn

694 + )

695

696 def forward(

697 self,

675 - hidden_states: torch.Tensor | None,

698 + pooled_data: TokensPoolingMethodOutputItem,

699 pooling_param: PoolingParams,

677 - ) -> PoolerOutput:

700 + ) -> TokensPoolerHeadOutput:

701 # for unfinished chunked prefill

679 - if hidden_states is None:

702 + if pooled_data is None:

703 return None

704

682 - hidden_states = hidden_states.to(self.head_dtype)

705 + pooled_data = pooled_data.to(self.head_dtype)

706 # hidden_states shape: [n_token, hidden_size]

707

708 if self.classifier is not None:

686 - scores = self.classifier(hidden_states)

709 + scores = self.classifier(pooled_data)

710 else:

688 - scores = hidden_states

711 + scores = pooled_data

712 # scores shape: [n_token, num_labels]

713

714 if self.logit_bias is not None:

715 scores -= self.logit_bias

716

717 if pooling_param.use_activation:

695 - scores = self.act_fn(scores)

718 + scores = self.activation(scores)

719

720 # scores shape: [n_token, num_labels]

721 return scores

722

723

724 class AllPooler(Pooler):

702 - def __init__(self, head: nn.Module | PoolerHead) -> None:

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

726 super().__init__()

727

728 self.pooling = AllPool()

@@ -712,17 +735,16 @@ class AllPooler(Pooler):

735 self,

736 hidden_states: torch.Tensor,

737 pooling_metadata: PoolingMetadata,

715 - ) -> PoolerOutput:

738 + ) -> TokensPoolerOutput:

739 pooled_data = self.pooling(hidden_states, pooling_metadata)

740 pooling_params = pooling_metadata.pooling_params

741 assert len(pooled_data) == len(pooling_params)

742

720 - pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]

721 - return pooled_data

743 + return [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]

744

745

746 class StepPooler(Pooler):

725 - def __init__(self, head: nn.Module | PoolerHead) -> None:

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

748 super().__init__()

749

750 self.pooling = AllPool()

@@ -730,14 +752,14 @@ class StepPooler(Pooler):

752

753 def extract_states(

754 self,

733 - hidden_states: torch.Tensor | list[torch.Tensor],

755 + hidden_states: torch.Tensor,

756 pooling_metadata: PoolingMetadata,

735 - ) -> PoolerOutput:

757 + ) -> list[torch.Tensor | None]:

758 pooled_data_lst = self.pooling(hidden_states, pooling_metadata)

759 prompt_token_ids = pooling_metadata.get_prompt_token_ids()

760 pooling_params = pooling_metadata.pooling_params

761

740 - pooled_data: PoolerOutput = []

762 + pooled_data = list[torch.Tensor | None]()

763 for data, token_id, pooling_param in zip(

764 pooled_data_lst, prompt_token_ids, pooling_params

765 ):

@@ -766,15 +788,14 @@ class StepPooler(Pooler):

788

789 def forward(

790 self,

769 - hidden_states: torch.Tensor | list[torch.Tensor],

791 + hidden_states: torch.Tensor,

792 pooling_metadata: PoolingMetadata,

771 - ) -> PoolerOutput:

793 + ) -> TokensPoolerOutput:

794 pooled_data = self.extract_states(hidden_states, pooling_metadata)

795 pooling_params = pooling_metadata.pooling_params

796 assert len(pooled_data) == len(pooling_params)

797

776 - pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]

777 - return pooled_data

798 + return [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]

799

800

801 class DispatchPooler(Pooler):

@@ -800,12 +821,12 @@ class DispatchPooler(Pooler):

821

822 def forward(

823 self,

803 - hidden_states: torch.Tensor | list[torch.Tensor],

824 + hidden_states: torch.Tensor,

825 pooling_metadata: PoolingMetadata,

826 ) -> PoolerOutput:

827 poolers_by_task = self.poolers_by_task

828

808 - outputs = list[torch.Tensor]()

829 + outputs = list[torch.Tensor | None]()

830 offset = 0

831 for task, group in groupby(pooling_metadata.tasks):

832 if not (pooler := poolers_by_task.get(task)):

vllm/model_executor/models/bert.py

+19 -14

@@ -24,11 +24,14 @@ from vllm.model_executor.layers.pooler import (

24 PoolingMethod,

25 PoolingParamsUpdate,

26 PoolingType,

27 + TokenPoolerHeadOutput,

28 + TokenPoolingMethodOutput,

29 )

30 from vllm.model_executor.layers.quantization import QuantizationConfig

31 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

32 from vllm.sequence import IntermediateTensors

33 from vllm.tasks import PoolingTask

34 +from vllm.v1.outputs import TokenPoolerOutput

35 from vllm.v1.pool.metadata import PoolingMetadata

37 from .interfaces import SupportsCrossEncoding, SupportsQuant

@@ -97,24 +100,26 @@ class BertPooler(Pooler):

100 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

101 return self.pooling.get_pooling_updates(task)

102

100 - def _head(self, pooled_output: torch.Tensor):

101 - pooled_output = self.dense(pooled_output)

102 - pooled_output = self.activation(pooled_output)

103 - return pooled_output

104 -

105 - def forward(

103 + def head(

104 self,

107 - hidden_states: torch.Tensor | list[torch.Tensor],

105 + pooled_data: TokenPoolingMethodOutput,

106 pooling_metadata: PoolingMetadata,

109 - ) -> torch.Tensor | list[torch.Tensor]:

110 - pooled_output = self.pooling(hidden_states, pooling_metadata)

107 + ) -> TokenPoolerHeadOutput:

108 + if isinstance(pooled_data, list):

109 + pooled_data = torch.stack(pooled_data)

110

112 - if isinstance(pooled_output, list):

113 - pooled_output = [self._head(output) for output in pooled_output]

114 - else:

115 - pooled_output = self._head(pooled_output)

111 + pooled_data = self.dense(pooled_data)

112 + pooled_data = self.activation(pooled_data)

113 + return pooled_data

114

117 - return pooled_output

115 + def forward(

116 + self,

117 + hidden_states: torch.Tensor,

118 + pooling_metadata: PoolingMetadata,

119 + ) -> TokenPoolerOutput:

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

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

122 + return pooled_data

123

124

125 class BertEncoder(nn.Module):

vllm/model_executor/models/gritlm.py

+17 -9

@@ -4,21 +4,22 @@ from collections.abc import Set

5 import numpy as np

6 import torch

7 -import torch.nn as nn

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

14 - PoolerHead,

13 PoolerNormalize,

14 + PoolingMethod,

15 PoolingParamsUpdate,

16 + TokenPoolerHeadOutput,

17 + TokenPoolingMethodOutput,

18 )

19 from vllm.model_executor.models.llama import LlamaForCausalLM

20 from vllm.tasks import PoolingTask

21 from vllm.tokenizers import cached_tokenizer_from_config

21 -from vllm.v1.outputs import PoolerOutput

22 +from vllm.v1.outputs import TokenPoolerOutput

23 from vllm.v1.pool.metadata import PoolingMetadata

25 from .interfaces_base import default_pooling_type

@@ -26,7 +27,7 @@ from .interfaces_base import default_pooling_type

27 logger = init_logger(__name__)

29 -class GritLMMeanPool(nn.Module):

30 +class GritLMMeanPool(PoolingMethod):

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

33 def __init__(self, model_config: ModelConfig):

@@ -141,16 +142,16 @@ class GritLMMeanPool(nn.Module):

142 return instruction_len

143

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

144 - return {"encode", "embed"}

145 + return {"embed"}

146

147 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

148 return PoolingParamsUpdate(requires_token_ids=True)

149

150 def forward(

151 self,

151 - hidden_states: torch.Tensor | list[torch.Tensor],

152 + hidden_states: torch.Tensor,

153 pooling_metadata: PoolingMetadata,

153 - ) -> list[torch.Tensor] | torch.Tensor:

154 + ) -> TokenPoolingMethodOutput:

155 prompt_lens = pooling_metadata.prompt_lens

156 instr_lens = torch.tensor(

157 [

@@ -178,7 +179,7 @@ class GritLMPooler(Pooler):

179 super().__init__()

180

181 self.pooling = GritLMMeanPool(model_config)

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

182 + self.activation = PoolerNormalize()

183

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

185 return self.pooling.get_supported_tasks()

@@ -186,11 +187,18 @@ class GritLMPooler(Pooler):

187 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

188 return self.pooling.get_pooling_updates(task)

189

190 + def head(

191 + self,

192 + pooled_data: TokenPoolingMethodOutput,

193 + pooling_metadata: PoolingMetadata,

194 + ) -> TokenPoolerHeadOutput:

195 + return self.activation(pooled_data)

196 +

197 def forward(

198 self,

199 hidden_states: torch.Tensor,

200 pooling_metadata: PoolingMetadata,

193 - ) -> PoolerOutput:

201 + ) -> TokenPoolerOutput:

202 pooled_data = self.pooling(hidden_states, pooling_metadata)

203 pooled_data = self.head(pooled_data, pooling_metadata)

204 return pooled_data

vllm/model_executor/models/modernbert.py

+18 -13

@@ -19,12 +19,15 @@ from vllm.model_executor.layers.pooler import (

19 PoolingMethod,

20 PoolingParamsUpdate,

21 PoolingType,

22 + TokenPoolerHeadOutput,

23 + TokenPoolingMethodOutput,

24 )

25 from vllm.model_executor.layers.rotary_embedding import get_rope

26 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding

27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

28 from vllm.sequence import IntermediateTensors

29 from vllm.tasks import PoolingTask

30 +from vllm.v1.outputs import TokenPoolerOutput

31 from vllm.v1.pool.metadata import PoolingMetadata

33 from .interfaces import SupportsCrossEncoding

@@ -300,23 +303,25 @@ class ModernBertPooler(Pooler):

303 def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:

304 return self.pooling.get_pooling_updates(task)

305

303 - def _head(self, pooled_output: torch.Tensor):

304 - pooled_output = pooled_output.to(self.dense.weight.dtype)

305 - return self.norm(self.act(self.dense(pooled_output)))

306 -

307 - def forward(

306 + def head(

307 self,

309 - hidden_states: torch.Tensor | list[torch.Tensor],

308 + pooled_data: TokenPoolingMethodOutput,

309 pooling_metadata: PoolingMetadata,

311 - ) -> torch.Tensor | list[torch.Tensor]:

312 - pooled_output = self.pooling(hidden_states, pooling_metadata)

310 + ) -> TokenPoolerHeadOutput:

311 + if isinstance(pooled_data, list):

312 + pooled_data = torch.stack(pooled_data)

313

314 - if isinstance(pooled_output, list):

315 - pooled_output = [self._head(output) for output in pooled_output]

316 - else:

317 - pooled_output = self._head(pooled_output)

314 + pooled_data = pooled_data.to(self.dense.weight.dtype)

315 + return self.norm(self.act(self.dense(pooled_data)))

316

319 - return pooled_output

317 + def forward(

318 + self,

319 + hidden_states: torch.Tensor,

320 + pooling_metadata: PoolingMetadata,

321 + ) -> TokenPoolerOutput:

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

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

324 + return pooled_data

325

326

327 @default_pooling_type("CLS")

vllm/v1/outputs.py

+4 -2

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

4 from abc import ABC, abstractmethod

5 from dataclasses import dataclass, field

6 -from typing import TYPE_CHECKING, NamedTuple

6 +from typing import TYPE_CHECKING, NamedTuple, TypeAlias

8 import numpy as np

9 import torch

@@ -91,7 +91,9 @@ class LogprobsTensors(NamedTuple):

92 # [num_reqs, <dynamic>]

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

94 -PoolerOutput = list[torch.Tensor | None] | torch.Tensor | None

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

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

96 +PoolerOutput: TypeAlias = TokenPoolerOutput | TokensPoolerOutput

99 @dataclass

vllm/v1/pool/metadata.py

@@ -90,6 +90,12 @@ class PoolingMetadata:

91 return [prompt_token_ids[i, :num] for i, num in enumerate(self.prompt_lens)]

93 + def get_pooling_cursor(self) -> PoolingCursor:

94 + pooling_cursor = self.pooling_cursor

95 + assert pooling_cursor is not None, "Should call `build_pooling_cursor` first"

96 +

97 + return pooling_cursor

98 +

99 def build_pooling_cursor(

100 self,

101 num_scheduled_tokens_np: np.ndarray,

vllm/v1/worker/gpu_model_runner.py

+1 -1

@@ -4680,7 +4680,7 @@ class GPUModelRunner(

4680 for task in supported_pooling_tasks:

4681 # Run a full batch with each task to ensure none of them OOMs

4682 output = self._dummy_pooler_run_task(hidden_states, task)

4683 - output_size[task] = sum(o.nbytes for o in output)

4683 + output_size[task] = sum(o.nbytes for o in output if o is not None)

4684 del output # Allow GC

4685

4686 max_task = max(output_size.items(), key=lambda x: x[1])[0]