Follow-up to #43286. Replaced the remaining ~10 assert statements across the pooling subsystem with proper exceptions to improve UX.
建议合并，这是一次低风险的质量改进，展示了如何用显式异常替代assert以提升代码健壮性和用户体验。值得对其他模块类似的assert进行类似替换。
review中noooop指出model-executor CI失败可能由本PR引起，作者随后更新测试以捕获RuntimeError，修复CI。无其他争议。
| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
vllm/model_executor/layers/pooler/seqwise/heads.py |
池化头部 | modified | 6.46 |
vllm/model_executor/layers/pooler/seqwise/methods.py |
池化方法 | modified | 6.18 |
vllm/config/pooler.py |
池化配置 | modified | 5.95 |
vllm/model_executor/layers/pooler/tokwise/heads.py |
Token 池化头 | modified | 5.94 |
vllm/model_executor/layers/pooler/seqwise/poolers.py |
序列池化器 | modified | 5.85 |
vllm/model_executor/layers/pooler/tokwise/poolers.py |
Token 池化器 | modified | 5.85 |
tests/model_executor/layers/test_pooler_methods.py |
池化测试 | modified | 4.06 |
vllm/model_executor/layers/pooler/seqwise/heads.py
data-contract
核心头部类，替换了两个 assert 为 ValueError，影响 embed 和 classify 流程
class EmbeddingPoolerHead(SequencePoolerHead):
def forward(
self,
pooled_data: SequencePoolingMethodOutput,
pooling_metadata: PoolingMetadata,
) -> SequencePoolerHeadOutput:
pooling_params = pooling_metadata.pooling_params
# 原 assert len(pooled_data) == len(pooling_params) 替换为 ValueError
if len(pooled_data) != len(pooling_params):
raise ValueError(
f"pooled_data length ({len(pooled_data)}) does not match "
f"pooling_params length ({len(pooling_params)})"
)
if isinstance(pooled_data, list):
pooled_data = torch.stack(pooled_data)
if self.head_dtype is not None:
pooled_data = pooled_data.to(self.head_dtype)
if self.projector is not None:
embeddings = self.projector(pooled_data)
else:
embeddings = pooled_data
# Matryoshka 维度截断
dimensions_list = [p.dimensions for p in pooling_params]
if any(d is not None for d in dimensions_list):
# 原 assert len(embeddings) == len(dimensions_list) 替换为 ValueError
if len(embeddings) != len(dimensions_list):
raise ValueError(
f"embeddings length ({len(embeddings)}) does not match "
f"dimensions_list length ({len(dimensions_list)})"
)
if len(set(dimensions_list)) == 1 and not isinstance(embeddings, list):
embeddings = embeddings[..., :dimensions_list[0]]
else:
embeddings = [
vecs if d is None else vecs[..., :d]
for vecs, d in zip(embeddings, dimensions_list)
]
# 归一化
if self.activation is not None:
flags = [p.use_activation for p in pooling_params]
if len(set(flags)) == 1:
if flags[0]:
embeddings = self.activation(embeddings)
else:
embeddings = [
self.activation(vecs) if f else vecs
for vecs, f in zip(embeddings, flags)
]
return embeddings
vllm/model_executor/layers/pooler/seqwise/methods.py
data-contract
CLSPool 和 MeanPool 中的 assert 替换为 RuntimeError，影响序列池化方法
class CLSPool(SequencePoolingMethod):
def forward(
self,
hidden_states: torch.Tensor,
pooling_metadata: PoolingMetadata,
) -> SequencePoolingMethodOutput:
pooling_cursor = pooling_metadata.get_pooling_cursor()
# 原 assert not pooling_cursor.is_partial_prefill() 替换为 RuntimeError
if pooling_cursor.is_partial_prefill():
raise RuntimeError("partial prefill is not supported with CLS pooling")
return hidden_states[pooling_cursor.first_token_indices_gpu]
class MeanPool(SequencePoolingMethod):
def forward(
self,
hidden_states: torch.Tensor,
pooling_metadata: PoolingMetadata,
) -> SequencePoolingMethodOutput:
pooling_cursor = pooling_metadata.get_pooling_cursor()
# 原 assert not pooling_cursor.is_partial_prefill() 替换为 RuntimeError
if pooling_cursor.is_partial_prefill():
raise RuntimeError("partial prefill is not supported with MEAN pooling")
prompt_lens_cpu = pooling_cursor.prompt_lens_cpu
num_seqs = prompt_lens_cpu.numel()
hidden_size = hidden_states.shape[-1]
if num_seqs == 0:
return hidden_states.new_empty((0, hidden_size), dtype=torch.float32)
# 剩余 chunked 计算不变 ...
vllm/config/pooler.py
core-logic
配置访问方法中的 assert 替换为 ValueError，避免未初始化时静默失败
def get_seq_pooling_type(self) -> SequencePoolingType:
# 原 assert self.seq_pooling_type is not None 替换为 ValueError
if self.seq_pooling_type is None:
raise ValueError(
"seq_pooling_type is not set; it should be resolved by"
" ModelConfig before calling get_seq_pooling_type()"
)
return self.seq_pooling_type
def get_tok_pooling_type(self) -> TokenPoolingType:
# 原 assert self.tok_pooling_type is not None 替换为 ValueError
if self.tok_pooling_type is None:
raise ValueError(
"tok_pooling_type is not set; it should be resolved by"
" ModelConfig before calling get_tok_pooling_type()"
)
return self.tok_pooling_type
noooop 指出 model-executor CI failure likely caused by this PR; 作者随后更新测试以接受 RuntimeError
结论：测试更新后 CI 修复 · 已解决
低风险。异常类型从AssertionError改为ValueError/RuntimeError，任何显式捕获AssertionError的代码可能需要更新，但在pooling模块外直接捕获的可能性很低。配置未初始化时现在抛出明确的ValueError而非AssertionError，行为更合理。
用户将看到更具描述性的错误信息（如'pooled_data length does not match pooling_params length'而非简单的AssertionError）。系统功能无变化，开发过程更健壮。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论