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 3, 08:00 UTC
[Misc] Aligning tokwise pooler heads for consistency (#43041)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
May 19, 2026 at 01:16 UTC
4a4fdabe28f3e2c8f9d05bcc80c4bf6d656b1ead
4 files changed
+9
-4
vllm/model_executor/layers/pooler/seqwise/poolers.py
+2
-2
@@ -89,7 +89,7 @@ class SequencePooler(Pooler):
89
return pooled_data
90
91
92
-def pooler_for_embed(pooler_config: PoolerConfig):
92
+def pooler_for_embed(pooler_config: PoolerConfig) -> SequencePooler:
93
pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
94
95
vllm_config = get_current_vllm_config()
@@ -109,7 +109,7 @@ def pooler_for_classify(
109
pooling: SequencePoolingMethod | SequencePoolingFn | None = None,
110
classifier: ClassifierFn | None = None,
111
act_fn: PoolerActivation | None = None,
112
-):
112
+) -> SequencePooler:
113
if pooling is None:
114
pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
115
vllm/model_executor/layers/pooler/tokwise/__init__.py
+4
@@ -18,6 +18,8 @@ from .methods import (
18
from .poolers import (
19
TokenPooler,
20
TokenPoolerOutput,
21
+ TokenPoolingFn,
22
+ TokenPoolingHeadFn,
23
pooler_for_token_classify,
24
pooler_for_token_embed,
25
)
@@ -34,6 +36,8 @@ __all__ = [
36
"get_tok_pooling_method",
37
"TokenPooler",
38
"TokenPoolerOutput",
39
+ "TokenPoolingFn",
40
+ "TokenPoolingHeadFn",
41
"pooler_for_token_classify",
42
"pooler_for_token_embed",
43
]
vllm/model_executor/layers/pooler/tokwise/heads.py
+2
-1
@@ -78,7 +78,8 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
78
# embeddings shape: [n_tokens, embedding_size]
79
80
# for matryoshka representation
81
- embeddings = embeddings[..., : pooling_param.dimensions]
81
+ if pooling_param.dimensions is not None:
82
+ embeddings = embeddings[..., : pooling_param.dimensions]
83
84
# for normalize
85
if self.activation is not None and pooling_param.use_activation:
vllm/model_executor/layers/pooler/tokwise/poolers.py
+1
-1
@@ -118,7 +118,7 @@ def pooler_for_token_classify(
118
pooling: TokenPoolingMethod | TokenPoolingFn | None = None,
119
classifier: ClassifierFn | None = None,
120
act_fn: PoolerActivation | None = None,
121
-):
121
+) -> TokenPooler:
122
if pooling is None:
123
pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
124