# PR #45210 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Misc][Model] add io processor for query/document embeddings from ColBERT (jinaai/jina-colbert-v2)
- 合并时间：2026-06-16 09:32
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/45210

---

# 执行摘要

- 一句话：为 ColBERT 添加 query/document 嵌入 IO 处理器
- 推荐动作：对于使用 ColBERT 模型的团队，强烈建议采用此插件以获取正确嵌入。对于其他开发者，可作为实现自定义 IOProcessor 插件的参考范例，展示了插件系统的使用方式、entry point 注册、类型定义和集成测试的模式。

# 功能与动机

根据 PR body，现有代码无法为 ColBERT 模型输出正确的 query 嵌入（应为 32x128）和 document 嵌入。该 PR 弥补了这一空白，确保输出与 Jina AI API 和原始 ColBERT 论文预期一致。关联 issue #42234。

# 实现拆解

1. **插件脚手架**：在 tests/plugins/colbert_query_plugin 下创建可安装 Python 包，通过 setup.py 注册 vllm.io_processor_plugins entry point，将 colbert_query_plugin 映射到 ColBERTQueryEmbeddingProcessor 类。
2. **数据模型定义**：types.py 定义了 InputType (query|document)、ColBERTEmbeddingCompletionRequestMixin（添加 input_type 字段）、ColBERTEmbeddingResponseData 和 ColBERTEmbeddingResponse，使用 Pydantic 模型。
3. **核心处理器**：ColBERTQueryEmbeddingProcessor 继承 IOProcessor，实现 parse_data 方法。parse_data 根据 input_type 分别调用 _build_query_prompt 或 _build_document_prompt 构建 TokensPrompt，然后执行推理，并将输出解析为 ColBERTEmbeddingResponse。
4. **提示构建**：_build_query_prompt 插入 [CLS]、[QueryMarker]、内容 token、[SEP]，然后填充 [MASK] 至 QUERY_MAXLEN=32；_build_document_prompt 仅插入 [CLS]、[DocumentMarker]、内容 token、[SEP]。
5. **CI 集成**：在 .buildkite/test-amd.yaml 和 .buildkite/test_areas/plugins.yaml 添加步骤，安装插件并运行集成测试 test_colbert_query_io_processor_plugins.py。

关键文件：
- `tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py`（模块 嵌入处理器；类别 source；类型 core-logic；符号 ColBERTQueryEmbeddingProcessor, __init__, __repr__, _resolve_marker_ids）: 核心插件实现，包含 ColBERTQueryEmbeddingProcessor 类，处理 query 和 document 提示构建及推理。
- `tests/plugins/colbert_query_plugin/colbert_query_processor/types.py`（模块 数据模型；类别 source；类型 data-contract；符号 ColBERTEmbeddingCompletionRequestMixin, ColBERTEmbeddingResponseData, ColBERTEmbeddingResponse）: 定义了输入输出数据模型，包括 input_type 字段和响应结构。
- `tests/plugins_tests/test_colbert_query_io_processor_plugins.py`（模块 集成测试；类别 test；类型 test-coverage；符号 ModelConfig, _get_attr_or_val, _check_token_embeddings, server）: 集成测试，验证 query/document 嵌入的正确性及错误处理。
- `tests/plugins/colbert_query_plugin/setup.py`（模块 打包脚本；类别 infra；类型 configuration）: 通过 entry_points 注册插件，使 vllm 能够发现 colbert_query_plugin。
- `tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py`（模块 插件入口；类别 infra；类型 entrypoint；符号 register_colbert_query_embedding_processor）: 暴露插件注册函数，返回处理器类的完全限定名。
- `.buildkite/test-amd.yaml`（模块 CI 配置；类别 infra；类型 configuration）: CI 配置中添加 colbert_query 插件的安装和测试步骤。
- `.buildkite/test_areas/plugins.yaml`（模块 CI 配置；类别 infra；类型 configuration）: 将 colbert_query 插件纳入插件测试区域。

关键符号：ColBERTQueryEmbeddingProcessor.parse_data, ColBERTQueryEmbeddingProcessor._build_query_prompt, ColBERTQueryEmbeddingProcessor._build_document_prompt, ColBERTQueryEmbeddingProcessor._resolve_marker_ids, ColBERTQueryEmbeddingProcessor._iter_content_token_ids, register_colbert_query_embedding_processor

## 关键源码片段

### `tests/plugins/colbert_query_plugin/colbert_query_processor/types.py`

定义了输入输出数据模型，包括 input_type 字段和响应结构。

```python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from typing import Literal, get_args

from pydantic import BaseModel, Field

from vllm.entrypoints.openai.engine.protocol import UsageInfo
from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin

# 输入类型：query 或 document
InputType = Literal["query", "document"]
INPUT_TYPES: tuple[InputType, ...] = get_args(InputType)
# query 嵌入向量的固定长度（ColBERT 论文中的 Q* 超参数）
QUERY_MAXLEN = 32


class ColBERTEmbeddingCompletionRequestMixin(CompletionRequestMixin):
    """请求扩展，增加 input_type 字段，描述输入应如何处理。"""
    input_type: InputType = Field(
        description="Whether to encode the input as a ColBERT 'query' "
        f"(query marker + [mask] expansion to {QUERY_MAXLEN} tokens) or as a "
        "'document' (document marker only). Required.",
    )


class ColBERTEmbeddingResponseData(BaseModel):
    """单个嵌入响应条目，包含 token 级别的嵌入向量列表。"""
    index: int
    object: str = "embedding"
    input_type: InputType
    embedding: list[list[float]]  # shape: (num_tokens, embedding_dim)


class ColBERTEmbeddingResponse(BaseModel):
    """嵌入响应顶层结构，兼容 vLLM 的 pooling 响应格式。"""
    data: list[ColBERTEmbeddingResponseData]
    usage: UsageInfo

```

# 评论区精华

Review 中 bennethinz 确认插件适用于 jinaai/jina-colbert-v2，但指出不适用于 colbert-ir/colbertv2.0 等变体。作者承诺在文档中明确说明支持范围。该讨论影响了插件的使用建议。

- 插件兼容性范围 (design): 作者同意并将在文档中明确说明只支持 jinaai/jina-colbert-v2。

# 风险与影响

- 风险：低风险。插件仅通过插件系统加载，不影响核心路径。主要风险包括：1）仅测试了 jinaai/jina-colbert-v2，其他 ColBERT 变体可能因 tokenizer 缺少特殊标记而失败；2）插件位于 tests/ 目录，可能被当作仅测试代码，但实际可通过 --io-processor-plugin 在生产中使用，需注意稳定性；3）CI 测试仅在 AMD 上运行，未覆盖 NVIDIA 环境。
- 影响：影响较小。仅使用 ColBERT 模型的用户受益，通过 /pooling 端点能正确获得 query/document 嵌入。对现有其他模型无影响。团队需注意该插件非官方模型支持，而是用户贡献的额外插件。
- 风险标记：插件特定模型限制 , 仅 AMD CI 测试 , 代码位于 tests 目录可能被忽视

# 关联脉络

- 暂无明显关联 PR