执行摘要

• 一句话：tokwise 与 seqwise pooler 一致性对齐
• 推荐动作：值得快速合并。该 PR 清理了小但影响一致性的技术债务，提升了代码质量和可维护性。开发者若使用 tokwise pooler 相关 API，可留意新增的导出符号。

功能与动机

根据 PR body，目的是解决 seqwise/tokwise 之间的一致性问题，包括缺失的 None 守卫、公共 API 导出不一致以及缺少返回类型注解。具体修复了 TokenEmbeddingPoolerHead 中 matryoshka dimensions 的 None 守卫，避免潜在切片错误；使 tokwise 公共 API 与 seqwise 对齐，导出 TokenPoolingFn 和 TokenPoolingHeadFn；并为三个 pooler_for_* 函数添加返回类型注解。

实现拆解

1. 修复 TokenEmbeddingPoolerHead 中 matryoshka dimensions 的 None 守卫（vllm/model_executor/layers/pooler/tokwise/heads.py）：在 forward_chunk 方法中，对 pooling_param.dimensions 增加 if pooling_param.dimensions is not None: 条件判断，避免当 dimensions 为 None 时切片操作 embeddings[..., : pooling_param.dimensions] 产生意外结果，与 seqwise 中已有的 None 守卫行为一致。

2. 导出 TokenPoolingFn 和 TokenPoolingHeadFn（vllm/model_executor/layers/pooler/tokwise/__init__.py）：在 __init__.py 中从 .poolers 模块导入 TokenPoolingFn 和 TokenPoolingHeadFn，并将它们加入 __all__ 列表，使 tokwise 公共 API 与 seqwise 对应导出对齐，便于外部调用。

3. 添加返回类型注解：为 pooler_for_embed（seqwise/poolers.py）、pooler_for_classify（seqwise/poolers.py）和 pooler_for_token_classify（tokwise/poolers.py）添加显式返回类型注解，分别标注为 -> SequencePooler 和 -> TokenPooler，提升类型安全性和代码可读性。

4. 测试配套：PR body 列出了多个已有的 pooling 测试套件（如 test_embedding.py、test_token_classification.py 等），作者确认所有测试应通过。本次变更未新增测试文件，依赖现有测试覆盖。

关键文件：

• vllm/model_executor/layers/pooler/tokwise/heads.py（模块 池化层；类别 source；类型 data-contract；符号 TokenEmbeddingPoolerHead）: 修复 matryoshka dimensions 的 None 守卫，避免潜在切片错误，与 seqwise 行为对齐。
• vllm/model_executor/layers/pooler/tokwise/__init__.py（模块 池化层；类别 source；类型 data-contract；符号 TokenPoolingFn, TokenPoolingHeadFn）: 导出 TokenPoolingFn 和 TokenPoolingHeadFn，使 tokwise 公共 API 与 seqwise 对齐。
• vllm/model_executor/layers/pooler/seqwise/poolers.py（模块 池化层；类别 source；类型 data-contract；符号 pooler_for_embed, pooler_for_classify）: 为 pooler_for_embed 和 pooler_for_classify 添加返回类型注解，提升类型安全性。
• vllm/model_executor/layers/pooler/tokwise/poolers.py（模块 池化层；类别 source；类型 data-contract；符号 pooler_for_token_classify）: 为 pooler_for_token_classify 添加返回类型注解。

关键符号：pooler_for_embed, pooler_for_classify, pooler_for_token_classify, TokenEmbeddingPoolerHead.forward_chunk

关键源码片段

vllm/model_executor/layers/pooler/tokwise/heads.py

修复 matryoshka dimensions 的 None 守卫，避免潜在切片错误，与 seqwise 行为对齐。

# vllm/model_executor/layers/pooler/tokwise/heads.py
# 修改前：直接切片，当 dimensions 为 None 时可能报错或产生错误形状
# 修改后：增加 None 守卫，与 seqwise 侧保持一致
​
def forward_chunk(self, pooled_data, pooling_param):
    # ...
    # for matryoshka representation
    if pooling_param.dimensions is not None: # 新增 None 检查
        embeddings = embeddings[..., : pooling_param.dimensions]
​
    # for normalize
    if self.activation is not None and pooling_param.use_activation:
        embeddings = self.activation(embeddings)
​
    return embeddings

vllm/model_executor/layers/pooler/tokwise/__init__.py

导出 TokenPoolingFn 和 TokenPoolingHeadFn，使 tokwise 公共 API 与 seqwise 对齐。

# vllm/model_executor/layers/pooler/tokwise/__init__.py
# 在导入部分新增 TokenPoolingFn 和 TokenPoolingHeadFn
from .poolers import (
    TokenPooler,
    TokenPoolerOutput,
    TokenPoolingFn, # 新增
    TokenPoolingHeadFn, # 新增
    pooler_for_token_classify,
    pooler_for_token_embed,
)
​
# 在 __all__ 中新增导出
__all__ = [
    # ... 原有符号 ...
    "TokenPoolingFn", # 新增
    "TokenPoolingHeadFn", # 新增
    # ... 其他符号 ...
]

vllm/model_executor/layers/pooler/seqwise/poolers.py

为 pooler_for_embed 和 pooler_for_classify 添加返回类型注解，提升类型安全性。

# vllm/model_executor/layers/pooler/seqwise/poolers.py
# 为 pooler_for_embed 添加返回类型注解
def pooler_for_embed(pooler_config: PoolerConfig) -> SequencePooler: # 新增返回类型
    # ... 原有实现 ...
​
# 为 pooler_for_classify 添加返回类型注解
def pooler_for_classify(
    pooler_config: PoolerConfig,
    *,
    pooling: SequencePoolingMethod | SequencePoolingFn | None = None,
    classifier: ClassifierFn | None = None,
    act_fn: PoolerActivation | None = None,
) -> SequencePooler: # 新增返回类型
    # ... 原有实现 ...

评论区精华

该 PR 没有人工 review 评论。gemini-code-assist[bot] 的自动审查确认了变更内容：提高了类型安全性、引入了 None 守卫并统一了导出。noooop 直接批准了 PR。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低。变更集中在类型注解、导出添加和 None 守卫修复，不涉及核心逻辑重构。pooling_param.dimensions is not None 的加入是防御性编程，不会破坏正常流程；类型注解仅在静态类型检查时生效；导出的新增不影响现有导入。潜在风险：若外部代码已通过非公开方式访问 TokenPoolingFn 或 TokenPoolingHeadFn，导出后可能产生命名冲突（可能性极低）。
• 影响：影响范围：仅影响 pooler 模块内部的类型一致性和 API 导出。影响程度：低。无功能变更，无性能影响。对于使用类型检查的工具（如 mypy、IDE）更友好；对于依赖 tokwise 公共 API 的开发者，现在可以直接使用 TokenPoolingFn 和 TokenPoolingHeadFn，与 seqwise 体验一致。
• 风险标记：无显著风险

关联脉络

• 暂无明显关联 PR