diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py
index 762869c6782..81e19001b76 100644
--- a/vllm/model_executor/layers/pooler/seqwise/poolers.py
+++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py
@@ -89,7 +89,7 @@ class SequencePooler(Pooler):
         return pooled_data
 
 
-def pooler_for_embed(pooler_config: PoolerConfig):
+def pooler_for_embed(pooler_config: PoolerConfig) -> SequencePooler:
     pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
 
     vllm_config = get_current_vllm_config()
@@ -109,7 +109,7 @@ def pooler_for_classify(
     pooling: SequencePoolingMethod | SequencePoolingFn | None = None,
     classifier: ClassifierFn | None = None,
     act_fn: PoolerActivation | None = None,
-):
+) -> SequencePooler:
     if pooling is None:
         pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
 
diff --git a/vllm/model_executor/layers/pooler/tokwise/__init__.py b/vllm/model_executor/layers/pooler/tokwise/__init__.py
index fbc610c8556..924c21fcf38 100644
--- a/vllm/model_executor/layers/pooler/tokwise/__init__.py
+++ b/vllm/model_executor/layers/pooler/tokwise/__init__.py
@@ -18,6 +18,8 @@ from .methods import (
 from .poolers import (
     TokenPooler,
     TokenPoolerOutput,
+    TokenPoolingFn,
+    TokenPoolingHeadFn,
     pooler_for_token_classify,
     pooler_for_token_embed,
 )
@@ -34,6 +36,8 @@ __all__ = [
     "get_tok_pooling_method",
     "TokenPooler",
     "TokenPoolerOutput",
+    "TokenPoolingFn",
+    "TokenPoolingHeadFn",
     "pooler_for_token_classify",
     "pooler_for_token_embed",
 ]
diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py
index 0377a86755a..d9f41132c06 100644
--- a/vllm/model_executor/layers/pooler/tokwise/heads.py
+++ b/vllm/model_executor/layers/pooler/tokwise/heads.py
@@ -78,7 +78,8 @@ class TokenEmbeddingPoolerHead(TokenPoolerHead):
         # embeddings shape: [n_tokens, embedding_size]
 
         # for matryoshka representation
-        embeddings = embeddings[..., : pooling_param.dimensions]
+        if pooling_param.dimensions is not None:
+            embeddings = embeddings[..., : pooling_param.dimensions]
 
         # for normalize
         if self.activation is not None and pooling_param.use_activation:
diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py
index 131074a5556..0b66097e381 100644
--- a/vllm/model_executor/layers/pooler/tokwise/poolers.py
+++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py
@@ -118,7 +118,7 @@ def pooler_for_token_classify(
     pooling: TokenPoolingMethod | TokenPoolingFn | None = None,
     classifier: ClassifierFn | None = None,
     act_fn: PoolerActivation | None = None,
-):
+) -> TokenPooler:
     if pooling is None:
         pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())