← 返回目录 模块 3 / 8

3.1 Generation Task

生成任务

GenerationTask 是 vLLM 中最核心、最广泛使用的任务类型,对应于标准的自回归语言模型(Autoregressive LM)推理,模型逐 Token 生成输出序列。在 vLLM 中,GenerationTask 包含三种具体任务:generate(文本生成)、transcription(语音转录)和 realtime(实时交互)。

从代码层面看,vLLM 的任务类型定义极为精简但语义明确:

# vllm/tasks.py (lines 5-28)
GenerationTask = Literal["generate", "transcription", "realtime"]
GENERATION_TASKS: tuple[GenerationTask, ...] = get_args(GenerationTask)

generate(文本生成)

generate 是 vLLM 最核心、最广泛使用的任务类型。它对应于标准的自回归语言模型推理,模型逐 Token 生成输出序列。

原理与机制:

  • 使用 PagedAttention 机制管理 KV Cache,将连续的 KV 缓存分页存储在非连续内存中。
  • 支持 Continuous Batching,在生成阶段动态批处理不同长度的请求。
  • 支持 Speculative Decoding、Prefix Caching、Chunked Prefill 等优化。

代码级细节:

  • 模型必须实现 VllmModelForTextGeneration 接口(定义于 vllm/model_executor/models/interfaces_base.py)。
  • 该接口要求模型实现 compute_logits(self, hidden_states) 方法,用于在生成阶段计算下一个 Token 的对数概率。
  • compute_logits 在 TP(Tensor Parallel)rank > 0 时返回 None,仅在 rank 0 上实际计算。

API 端点:

  • 离线:LLM.generate(prompts, sampling_params)
  • 在线:POST /v1/completions(兼容 OpenAI API)、POST /v1/chat/completions

transcription(语音转录)

transcription 任务专用于 Whisper 等自动语音识别(ASR)模型。它将音频输入转录为文本输出。

原理与机制:

  • 基于 Encoder-Decoder 架构(如 Whisper)。
  • 音频通过编码器转换为隐藏状态,解码器自回归生成转录文本。
  • vLLM 对 Whisper 的支持包括 WhisperForConditionalGeneration 等架构。

代码级细节:

  • 在模型注册表中,WhisperForConditionalGeneration 被归类为 _MULTIMODAL_MODELS。
  • 模型通过 supports_transcription 接口标记声明支持转录任务。
  • 音频输入通过多模态预处理管道转换为 mel-spectrogram 特征。

realtime(实时交互)

realtime 任务是 vLLM 较新引入的任务类型,用于支持实时音频/语音交互场景。这对应于 OpenAI 的 Realtime API 模式,模型需要以低延迟处理流式音频输入并产生流式音频或文本输出。

原理与机制:

  • 支持流式输入和输出。
  • 需要模型具备低延迟处理能力。
  • 通常与 VAD(Voice Activity Detection)和音频 tokenization 结合。

代码级细节:

  • 在注册表中,Qwen3ASRRealtimeGeneration 和 VoxtralRealtimeGeneration 支持此任务。
  • 模型需要实现特定的实时处理接口。
  • 调度器需要支持更细粒度的请求切分和优先级管理。

Generation Task 与其他概念的关系

Generation Task 依赖 PagedAttention 和 Continuous Batching:所有生成任务的核心执行路径都使用 PagedAttention 管理 KV Cache,并通过 Continuous Batching 进行调度。没有这两项技术,vLLM 无法高效地服务大量并发生成请求。

Generation Task 使用 Sampling 策略:生成任务在每一步都需要从模型输出的 logits 分布中采样下一个 token。vLLM 支持 Greedy、Temperature、Top-p、Top-k、Beam Search 等多种采样策略(详见模块四 4.7 节)。

3.2 Pooling Task

池化任务

PoolingTask 是 vLLM 对非生成类模型的统一抽象,涵盖了嵌入、分类、Token 级任务等。这些模型的共同特点是:它们不对隐藏状态进行自回归解码,而是通过某种 "pooling" 操作从隐藏状态中提取固定维度的输出。

从代码层面:

PoolingTask = Literal[
    "embed",
    "classify",
    "token_embed",
    "token_classify",
    "plugin",
    "embed&token_classify",
]
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)

embed(序列嵌入)

embed 任务将输入序列(文本、图像等)编码为单个固定维度的向量表示。这是最常见的嵌入模型任务。

原理与机制:

  • 输入经过 Transformer 编码器产生每个 Token 的隐藏状态。
  • 通过 Sequence Pooling(如 LAST、MEAN、CLS)将 Token 级隐藏状态聚合为单个序列级向量。
  • 可选归一化(L2 normalize)使向量位于单位球面上。

代码级细节:

  • 对应 PoolerConfig.task = "embed"。
  • 默认 seq_pooling_type = "LAST",use_activation = True(归一化)。
  • 输出通过 EmbeddingPoolerHead 处理,支持 Matryoshka 降维和 Sentence Transformers projector。

API 端点:

  • 离线:LLM.embed(...)、LLM.encode(..., pooling_task="embed")
  • 在线:POST /v1/embeddings(OpenAI 兼容)、POST /v2/embed(Cohere 兼容)、POST /pooling

classify(序列分类)

classify 任务对输入序列进行分类,输出每个预定义类别的概率分布。

原理与机制:

  • 输入经过编码器后,通过 Sequence Pooling 提取序列表示。
  • 序列表示经过分类头(通常是线性层)映射到类别数量维度。
  • 应用 softmax(或 sigmoid,对于多标签)得到概率分布。

代码级细节:

  • 对应 PoolerConfig.task = "classify"。
  • 默认 seq_pooling_type = "LAST",use_activation = True(softmax)。
  • 支持 logit_mean 和 logit_sigma 进行仿射分数校准(Platt scaling)。
  • 输出通过 ClassifierPoolerHead 处理。

token_embed(Token 嵌入)

token_embed 任务保留每个输入 Token 的向量表示,不进行序列级聚合。这支持需要细粒度 Token 级表示的场景。

原理与机制:

  • 输入经过编码器后,直接输出每个 Token 的隐藏状态(或经过投影后的表示)。
  • 不进行 CLS/MEAN/LAST 等序列级 pooling。
  • 支持 late-interaction(如 ColBERT 系列模型)。

代码级细节:

  • 对应 PoolerConfig.task = "token_embed"。
  • 默认 tok_pooling_type = "ALL"。
  • 输出通过 TokenEmbeddingPoolerHead 处理。
  • 支持 chunked prefill 时的隐藏状态缓存和拼接。

token_classify(Token 分类)

token_classify 任务对每个输入 Token 进行分类,输出每个 Token 的类别概率分布。这用于命名实体识别(NER)、分词、词性标注、奖励模型等场景。

原理与机制:

  • 输入经过编码器后,每个 Token 的隐藏状态独立经过分类头。
  • 输出形状为 [batch_size, seq_len, num_labels]。
  • 支持 step_tag_id 过滤(用于奖励模型中的步骤标签)。

代码级细节:

  • 对应 PoolerConfig.task = "token_classify"。
  • 默认 tok_pooling_type = "ALL"。
  • 输出通过 TokenClassifierPoolerHead 处理。
  • 支持 returned_token_ids 提取特定维度(如奖励模型中的 good/bad token)。

plugin(插件任务)

plugin 任务是 vLLM 为 IO Processor 插件预留的任务类型。它允许用户自定义处理流程,在模型推理前后插入自定义逻辑。

embed&token_classify(组合任务)

embed&token_classify 是组合任务,允许模型同时输出序列嵌入和 Token 级分类结果。这在某些多任务学习模型中有用。

Pooling Task 与 Score Type 的映射

vLLM 通过 SCORE_TYPE_MAP 将 PoolingTask 映射到 ScoreType:

PoolingTaskScoreType用途
embedbi-encoder独立编码 query 和 document,用余弦相似度评分
classifycross-encoder联合编码 query+document,用线性分类器评分
token_embedlate-interaction保留 Token 级表示,用 MaxSim 评分

这一映射决定了 LLM.score(...) API 使用哪种评分机制。

3.3 Embedding Model

嵌入模型

Embedding Model(嵌入模型)是一类将离散的非结构化数据(文本、图像、音频等)映射到连续向量空间(通常是 Rd)的机器学习模型。在 vLLM 中,嵌入模型属于 Pooling Model 的子集,通过 embed 任务提供服务。

从信息检索的角度,嵌入模型通常采用 Bi-encoder(双编码器) 架构:query 和 document 分别独立编码,然后通过向量相似度(如余弦相似度)计算相关性。

Bi-encoder 架构的详细工作原理

架构定义:

Bi-encoder 由两个共享权重的编码器组成(通常实际上是同一个模型):

q_vec = f_Q(query), d_vec = f_D(document)

其中 f_Q 和 f_D 通常是相同的 Transformer 编码器(如 BERT、RoBERTa、LLaMA 等),q_vec, d_vec in Rd 是输出的嵌入向量。

相关性得分通过向量相似度计算:

score(q, d) = (q_vec . d_vec) / (||q_vec|| ||d_vec||) = cos(theta)

即余弦相似度,取值范围为 [-1, 1]。在实际应用中,通常会先对嵌入向量进行 L2 归一化,这样内积等价于余弦相似度。

独立编码的优势与代价:

优势:

  1. 可预计算:document 嵌入可以离线预计算并存储在向量数据库(如 Milvus、Pinecone、Qdrant、Weaviate)中。
  2. 高吞吐:query 时只需编码 query,然后在向量索引中做近似最近邻(ANN)搜索,毫秒级响应。
  3. 可扩展:可以轻松扩展到数十亿文档的索引。

代价:

  1. 精度损失:query 和 document 在编码时没有交互,无法捕捉细粒度的词级匹配关系。
  2. 信息瓶颈:将整个文档压缩为单个向量,长文档的信息可能丢失。
  3. 语义偏差:不同意图的文档可能有相似的嵌入。

Pooling 方式:CLS、LAST、MEAN

在 Bi-encoder 中,将 Token 级隐藏状态聚合为序列级向量是核心步骤。vLLM 支持三种序列级 Pooling 方式:

CLS Pooling:

h_CLS = H[0]

其中 H in Rn x d 是最后一层 Transformer 的隐藏状态矩阵,H[0] 取第一个位置([CLS] token)的隐藏状态。

适用场景:BERT、RoBERTa 等双向编码器模型。原始 BERT 论文中,[CLS] 被设计为聚合序列信息的特殊 token。

注意:CLS pooling 不支持 chunked prefill,因为需要完整的序列才能确定 CLS 位置。

LAST Pooling:

h_LAST = H[n-1]

取序列最后一个有效 token 的隐藏状态。对于因果(causal)语言模型(如 LLaMA、GPT、Qwen),最后一个 token 的隐藏状态天然聚合了前面所有 token 的信息(由于自回归掩码)。

适用场景:因果语言模型作为嵌入模型时(如 e5-mistral、GritLM)。vLLM 的默认 pooling 方式。

LAST pooling 天然支持 chunked prefill,因为只需要知道最后一个 token 的索引。

MEAN Pooling:

h_MEAN = (1/n) * sum_{i=0}^{n-1} H[i]

对所有有效 token 的隐藏状态取平均。这比 LAST pooling 更充分利用序列中所有 token 的信息。

适用场景:Sentence Transformers 系列模型(如 all-MiniLM-L6-v2、multilingual-e5-large)。

注意 MEAN pooling 也不支持 chunked prefill(需要完整序列),但 vLLM 通过 enable_chunked_processing 参数为长文本提供了分块处理机制。

Matryoshka Representation Learning(嵌套表示学习)

Matryoshka Representation Learning(MRL,嵌套表示学习/俄罗斯套娃表示学习)是一种训练技术,使得嵌入向量的前 m 个维度(m < d)本身构成一个有效的 m 维嵌入,无需额外的降维操作。

核心思想:

  • 传统嵌入模型训练时只优化完整维度的损失。
  • MRL 在训练时对多个截断维度同时计算损失并求和。
  • 例如,对于 768 维嵌入,同时优化 768、512、256、128、64 维的损失。

数学形式:

L_MRL = sum_{m in M} w_m * L(f(x)_{1:m}, f(x')_{1:m})

其中 M 是预定义的维度集合(如 {64, 128, 256, 512, 768}),w_m 是各维度的权重,f(x)_{1:m} 表示取前 m 维。

为什么 MRL 有效:

MRL 的关键洞察是:嵌入向量的不同维度承载的信息重要性不同。通过联合优化,模型被迫将最重要的语义信息编码到前面的维度中,次要信息放在后面的维度。

实验表明,MRL 嵌入即使截断到 8.3% 的原始尺寸,仍能保留 98.37% 的性能。这远优于后处理降维方法(如 PCA、SVD),因为:

  1. PCA 是数据相关的,需要拟合转换矩阵。
  2. MRL 的降维只需简单截断,零额外计算开销。
  3. MRL 的优化目标直接面向下游检索质量。

vLLM 中的 MRL 支持:

vLLM 在 PoolerConfig 中提供 dimensions 参数支持 MRL:

class PoolerConfig:
    dimensions: int | None = None
    """Reduce the dimensions of embeddings if model support matryoshka
    representation. Defaults to None."""

在 EmbeddingPoolerHead 中的实现:

class EmbeddingPoolerHead:
    def forward(self, pooled_data, pooling_metadata):
        dimensions_list = [pooling_param.dimensions for pooling_param in pooling_params]
        if any(d is not None for d in dimensions_list):
            if len(set(dimensions_list)) == 1 and not isinstance(embeddings, list):
                d = dimensions_list[0]
                embeddings = embeddings[..., :d]  # 截断到前 d 维

注意事项:并非所有嵌入模型都支持 MRL。vLLM 会检查模型配置中的 is_matryoshka 或 matryoshka_dimensions 字段。对于不支持的模型设置 dimensions 参数会报错。

常见 Embedding 模型

BertModel / RobertaModel:基于 Transformer 编码器,双向注意力。原始 BERT-base 有 110M 参数,hidden_size=768。使用 [CLS] token 进行序列表示(配合 CLS pooling)。典型模型:BAAI/bge-base-en-v1.5、sentence-transformers/all-MiniLM-L6-v2、BAAI/bge-m3。

GteModel / SnowflakeGteNewModel:GTE(General Text Embeddings)是阿里巴巴达摩院开发的嵌入模型系列。基于 BERT-with-RoPE 架构,结合 RoPE(旋转位置编码)和双向注意力。

JinaEmbeddings:Jina AI 开发的嵌入模型。以其优秀的多语言支持和长文本处理能力著称。

ColBERT(作为 Token Embedding 模型):虽然 ColBERT 主要归类为 late-interaction 模型,但它在 token_embed 任务中也产生 Token 级嵌入。

在 RAG 流程中的位置和作用

在典型的 RAG(Retrieval-Augmented Generation)流程中,Embedding Model 位于 检索阶段(Retrieval Stage):

User Query -> Query Embedding -> Vector Search (ANN) -> Top-K Documents
                                                  ^
                                    Document Embeddings (Precomputed & Indexed)

具体作用:

  1. 索引阶段:将知识库中的文档分块(chunking),用 Embedding Model 编码为向量,存入向量数据库。
  2. 查询阶段:将用户查询编码为向量,在向量数据库中检索最相似的文档块。
  3. 增强生成:将检索到的文档块与用户查询一起送入生成模型(LLM)产生回答。

性能考量:

  • Bi-encoder 的检索精度通常低于 cross-encoder,但速度是 cross-encoder 的 100-1000 倍。
  • 典型 RAG 采用 "Bi-encoder 粗排 + Cross-encoder/Reranker 精排" 的两阶段策略。
  • 在 MS MARCO 等基准上,Bi-encoder Top-100 精度约 60%,经 Cross-encoder 重排后可达 85%+。

3.4 Rerank Model

重排序模型

Rerank Model(重排序模型)是 RAG 流程中 精排阶段(Reranking Stage) 使用的模型。它接收一个 query 和候选文档,输出它们的相关性分数,用于对粗排结果进行重新排序。

在 vLLM 中,reranker 是 classification model 的一个子集:当分类模型的 num_labels == 1 时,它可以作为评分模型(scoring model)使用,其输出经过适当的激活函数后即为相关性分数。

Cross-encoder 架构的详细工作原理

架构定义:

Cross-encoder 与 Bi-encoder 的根本区别在于:query 和 document 不是独立编码,而是被拼接成单个序列后一起输入 Transformer:

Input: [CLS] query tokens [SEP] document tokens [SEP]
score(q, d) = f_cls(Transformer([q; d]))

其中 f_cls 通常是 [CLS] token 的隐藏状态经过一个线性层得到的标量分数。

全交互注意力机制:

Cross-encoder 的核心优势在于 early interaction(早期交互):

  • Query 的每个 token 都可以通过自注意力机制直接 attend to document 的每个 token。
  • 模型可以捕捉细粒度的词级匹配关系(如 "machine learning" 与 "deep learning" 的局部语义关联)。
  • 可以识别否定、条件等复杂语义关系。

数学形式:

H = Transformer([q_1, ..., q_m, d_1, ..., d_n])
score = W . H_[CLS] + b

其中 W in R1 x d,b in R 是分类头的参数。

与 Bi-encoder 的详细对比

维度Bi-encoderCross-encoder
架构独立双编码器单编码器联合处理
编码方式Query 和 Doc 分别编码Query 和 Doc 拼接编码
交互时机无交互 / 晚期交互早期交互
相似度计算余弦相似度 / 点积线性分类器输出
速度快(ms 级,支持 ANN)慢(每对需要一次前向传播)
可预计算Document 可离线预计算不可预计算
精度中等(NDCG@10 ~ 30-40)高(NDCG@10 ~ 45-55)
适用阶段粗排(召回)精排(重排序)
典型延迟5ms / 百万文档50ms / 100 文档
存储成本低(每文档一个向量)无额外存储
长文本处理受限于最大长度,需截断同样受限,但交互更细粒度

实际应用场景:

  • Bi-encoder:从百万级文档库中快速召回 Top-100 候选。
  • Cross-encoder:对这 100 个候选重新排序,选出 Top-10 最相关的。

在 RAG 流程中的精排阶段

典型的两阶段 RAG 检索流程:

Stage 1 (Retrieval - Bi-encoder):
  Query -> Embed -> ANN Search -> Top-K Candidates (K=100~1000)

Stage 2 (Reranking - Cross-encoder):
  (Query, Doc_i) for i in Top-K -> Cross-encode -> Score -> Sort -> Top-M (M=5~20)

Stage 3 (Generation - LLM):
  Query + Top-M Documents -> LLM -> Answer

Cross-encoder 在精排阶段的提升效果显著:

  • Bi-encoder Top-10 精度约 60%
  • 经 Cross-encoder 重排后可达 85%+
  • 提升主要来自纠正 Bi-encoder 的 "假阳性"(语义相似但不相关)和 "假阴性"(语义不相似但相关)

vLLM 中常见的 Rerank 模型和实现

基于序列分类的 Reranker:

在模型注册表中,_SEQUENCE_CLASSIFICATION_MODELS 中的模型可用于 reranking:

_SEQUENCE_CLASSIFICATION_MODELS = {
    "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
    "GteNewForSequenceClassification": ("bert_with_rope", "GteNewForSequenceClassification"),
    "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"),
    "ModernBertForSequenceClassification": ("modernbert", "ModernBertForSequenceClassification"),
    "LlamaBidirectionalForSequenceClassification": ("llama", "LlamaBidirectionalForSequenceClassification"),
    # ...
}

这些模型通过 as_seq_cls_model 适配器将生成模型转换为分类模型。关键代码在 adapters.py:

def as_seq_cls_model(cls):
    class ModelForSequenceClassification(_create_pooling_model_cls(cls), SupportsCrossEncoding):
        def _init_pooler(self, vllm_config, prefix=""):
            self.score = ReplicatedLinear(
                model_config.get_hidden_size(),
                text_config.num_labels,
                bias=False,
                params_dtype=model_config.head_dtype,
                quant_config=quant_config,
            )
            return DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)

注意 SupportsCrossEncoding 接口标记了 score_type = "cross-encoder",这使得 Score API 知道使用 classify 任务而非 embed 任务。

Jina Reranker:

"JinaForRanking": ("jina", "JinaForRanking"),

Jina Reranker 是基于 Qwen3 架构的专门重排序模型,在多个基准上表现优异。

Qwen3 Reranker:

vllm serve Qwen/Qwen3-Reranker-0.6B \
  --hf_overrides '{"architectures": ["Qwen3ForSequenceClassification"],"classifier_from_token": ["no", "yes"],"is_original_qwen3_reranker": true}'

Rerank Model 与 Embedding Model 的关系

Rerank Model 是 Embedding Model 的精排补充。在 RAG 流程中:

  • Embedding Model(Bi-encoder)负责从海量文档中快速召回候选(粗排)。
  • Rerank Model(Cross-encoder)负责对候选进行精细排序(精排)。
  • 两者是串联关系:没有 Bi-encoder 的粗排,Cross-encoder 无法处理百万级文档;没有 Cross-encoder 的精排,Bi-encoder 的召回质量不足以直接用于生成。

3.5 Classification Model

分类模型

Classification Model(分类模型)将输入序列映射到预定义类别集合的概率分布。在 vLLM 中,分类模型通过 classify 和 token_classify 任务提供服务,是 _SEQUENCE_CLASSIFICATION_MODELS 和 _TOKEN_CLASSIFICATION_MODELS 的核心用途。

序列分类(Sequence Classification)

定义:为整个输入序列输出一个类别概率分布。

输出形状:[batch_size, num_labels]

典型应用:

  • 情感分析(正面/负面/中性)
  • 主题分类(体育/政治/科技)
  • 文本蕴含(蕴含/矛盾/中立)
  • Reranking(相关/不相关,num_labels=1 时)

vLLM 中的实现:

  • 使用 SequencePooler + ClassifierPoolerHead
  • 默认 seq_pooling_type = "LAST"
  • 默认激活函数:softmax(单标签)或 sigmoid(多标签)

代码级细节(seqwise/heads.py 中的 ClassifierPoolerHead):

class ClassifierPoolerHead(SequencePoolerHead):
    def __init__(self, classifier, logit_mean, logit_sigma, head_dtype, activation):
        self.classifier = classifier  # 线性分类层
        self.logit_mean = logit_mean  # 分数校准:均值偏移
        self.logit_sigma = logit_sigma  # 分数校准:尺度缩放
        self.activation = activation  # 如 softmax

    def forward(self, pooled_data, pooling_metadata):
        if self.classifier is not None:
            logits = self.classifier(pooled_data)  # [batch, num_labels]

        # Affine score calibration
        if self.logit_mean is not None:
            logits = logits - self.logit_mean
        if self.logit_sigma is not None:
            logits = logits / self.logit_sigma

        if self.activation is not None:
            logits = self.activation(logits)

        return logits

Token 分类(Token Classification / NER)

定义:为输入序列中的每个 Token 输出一个类别概率分布。

输出形状:[batch_size, seq_len, num_labels]

典型应用:

  • 命名实体识别(NER):识别 Person、Organization、Location 等
  • 词性标注(POS)
  • 强制对齐(Forced Alignment,如 ASR 中的音素对齐)
  • 奖励模型(评估每个生成步骤的质量)

vLLM 中的实现:

  • 使用 TokenPooler + TokenClassifierPoolerHead
  • 默认 tok_pooling_type = "ALL"
  • 支持 step_tag_id 过滤(奖励模型专用)
  • 支持 returned_token_ids 提取特定维度

代码级细节(tokwise/heads.py 中的 TokenClassifierPoolerHead):

class TokenClassifierPoolerHead(TokenPoolerHead):
    def forward_chunk(self, pooled_data, pooling_param):
        if self.classifier is not None:
            logits = self.classifier(pooled_data)  # [n_tokens, num_labels]

        if self.logit_mean is not None:
            logits = logits - self.logit_mean
        if self.logit_sigma is not None:
            logits = logits / self.logit_sigma

        if self.activation is not None and pooling_param.use_activation:
            logits = self.activation(logits)

        return logits

logit_mean / logit_sigma 分数校准

logit_mean 和 logit_sigma 是 vLLM 提供的 仿射分数校准(Affine Score Calibration) 参数,对应于 Platt Scaling 的简化形式:

p = sigma( (logit - mu) / sigma_scale )

其中:

  • mu = logit_mean:logits 的均值偏移量
  • sigma_scale = logit_sigma:logits 的尺度缩放因子
  • sigma(.) 是激活函数(通常是 softmax 或 sigmoid)

为什么需要分数校准:

神经网络输出的原始 logits 分布可能:

  1. 有偏(Biased):均值不为 0,导致概率估计系统性地偏高或偏低。
  2. 尺度不当:方差过大或过小,导致概率过于极端(接近 0 或 1)或过于平缓。

分数校准通过简单的仿射变换 z = (x - mu) / sigma 将 logits 调整到更合适的范围,使得输出的概率更接近真实的置信度。

在 PoolerConfig 中:

class PoolerConfig:
    logit_mean: float | None = None
    """If provided, subtract this value from classification logits before activation."""

    logit_sigma: float | None = None
    """If provided, divide the classification logits by this value after mean subtraction."""

历史兼容性:logit_bias(已弃用,v0.21 移除)映射到 logit_mean;logit_scale(已弃用,v0.21 移除)映射到 logit_sigma = 1/logit_scale。

3.6 Reward Model

奖励模型

Reward Model(奖励模型)是 RLHF(Reinforcement Learning from Human Feedback)流程中的核心组件。它接收一个 prompt-completion 对,输出一个标量奖励值,用于评估生成质量,替代人类直接标注偏好。

在 vLLM 中,奖励模型通过 token_classify 任务实现,因为奖励模型本质上是 Token 级分类器:它为每个生成步骤(或最终输出)预测一个质量分数。

在 RLHF 流程中的位置和作用

RLHF 三阶段流程:

Stage 1: SFT (Supervised Fine-Tuning)
  预训练模型 + 高质量指令数据 -> 指令遵循模型

Stage 2: Reward Modeling
  指令遵循模型 + 人类偏好比较数据 -> 奖励模型
  数据格式:(prompt, completion_A, completion_B, preference)
  训练目标: Bradley-Terry 模型,最大化偏好数据的似然

Stage 3: RL Optimization (PPO / DPO / IPO / etc.)
  指令遵循模型 + 奖励模型 -> 对齐人类偏好的最终模型
  目标:最大化奖励,同时保持与 SFT 模型的 KL 散度约束

奖励模型的具体作用:

  1. 偏好压缩:将复杂的人类偏好(如有用性、无害性、诚实性)压缩为标量奖励。
  2. 可扩展的评估:训练后无需人类即可自动评估任意生成结果。
  3. RL 优化信号:为策略梯度方法(如 PPO)提供可微分的优化目标。

数学形式(Bradley-Terry 模型):

P(y_w > y_l | x) = sigma( r(x, y_w) - r(x, y_l) )

其中 r(x, y) 是奖励模型对 prompt x 和 completion y 的输出,y_w 是人类偏好的输出,y_l 是非偏好的输出,sigma 是 sigmoid 函数。

step_tag_id 和奖励模型步长标签

step_tag_id 是 vLLM PoolerConfig 中的一个关键参数,用于 Process Reward Models(PRM,过程奖励模型):

class PoolerConfig:
    step_tag_id: int | None = None
    """If set, only the score corresponding to the step_tag_id in the
    generated sentence should be returned."""

    returned_token_ids: list[int] | None = None
    """A list of indices for the vocabulary dimensions to be extracted,
    such as the token IDs of good_token and bad_token."""

过程奖励模型(PRM):

PRM 不仅评估最终答案,还评估中间推理步骤的质量。这在数学推理、代码生成等需要多步思考的任务中尤为重要。

工作原理:

  • 模型在生成过程中插入特殊的步骤标签(如 <step>)。
  • step_tag_id 是这个特殊标签对应的 vocabulary ID。
  • Pooler 在 StepPool 中识别这些标签的位置,只提取步骤标签处的隐藏状态作为奖励。

代码实现(tokwise/methods.py 中的 StepPool):

class StepPool(AllPool):
    def forward(self, hidden_states, pooling_metadata):
        pooled_data_lst = super().forward(hidden_states, pooling_metadata)
        prompt_token_ids_cpu = pooling_metadata.get_prompt_token_ids_cpu()

        for data, token_id_cpu, pooling_param in zip(pooled_data_lst, ...):
            step_tag_id = pooling_param.step_tag_id
            returned_token_ids = pooling_param.returned_token_ids

            if returned_token_ids is not None:
                data = data[:, returned_token_ids]  # 只保留指定维度

            if step_tag_id is not None:
                idx_cpu = (token_id_cpu == step_tag_id).nonzero(as_tuple=True)[0]
                idx = idx_cpu.to(data.device, non_blocking=True)
                data = data[idx]  # 只保留 step_tag 位置的输出

            pooled_data.append(data)

常见 Reward 模型

InternLM2ForRewardModel:

_REWARD_MODELS = {
    "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
}

典型模型:internlm/internlm2-1_8b-reward、internlm/internlm2-7b-reward。基于 InternLM2 架构,支持 LoRA 和 Pipeline Parallelism,输出单维度奖励分数。

Qwen2ForRewardModel / Qwen2ForProcessRewardModel:

_REWARD_MODELS = {
    "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
    "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
}

典型模型:Qwen/Qwen2.5-Math-RM-72B、Qwen/Qwen2.5-Math-PRM-7B。

3.7 Score Type 体系

评分类型

vLLM 的 Score Type 体系定义了三种不同的语义相似度计算范式,分别对应不同的模型架构和应用场景。Score Type 在 vllm/tasks.py 中定义为:

ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
    "embed": "bi-encoder",
    "classify": "cross-encoder",
    "token_embed": "late-interaction",
}

bi-encoder(余弦相似度)

Bi-encoder 独立编码 query 和 document,然后通过向量相似度评分。

数学公式:

score(q, d) = (q_vec . d_vec) / (||q_vec|| ||d_vec||) = cos(theta)

工作流程:

  1. 分别编码 query 和 document 为向量 q_vec 和 d_vec。
  2. 计算余弦相似度(或点积)。
  3. 得分越高表示越相关。

特点:可预计算 document 嵌入;速度快,支持 ANN 搜索;但 query 和 document 在编码时无交互,精度有限。

cross-encoder(线性分类器输出)

Cross-encoder 将 query 和 document 拼接后联合编码,输出标量相关性分数。

数学公式:

score(q, d) = W . H_[CLS] + b

其中 H = Transformer([q; d]),W in R1 x d,b in R。

工作流程:

  1. 将 query 和 document 拼接为单个序列。
  2. 输入 Transformer 编码器。
  3. 取 [CLS] 隐藏状态经过线性层得到分数。

特点:query 和 document 在编码时充分交互,精度高;不可预计算,每对需要一次前向传播;速度慢,适合精排阶段。

late-interaction(MaxSim,ColBERT 系列)

Late-interaction 是一种折中方案:独立编码 query 和 document 到 token 级向量,但在最后阶段进行细粒度交互(MaxSim)。

数学公式:

对于 query token 嵌入 Q in Rm x d 和 document token 嵌入 D in Rn x d:

score(Q, D) = sum_{i=1}^{m} max_{j=1}^{n} (Q_i . D_j)

即:对每个 query token,找到与之最相似的 document token,取点积;然后对所有 query token 的 best match 求和。

工作流程:

  1. 分别编码 query 和 document 为 token 级向量矩阵。
  2. 对每个 query token,计算与所有 document token 的点积,取最大值。
  3. 对所有 query token 的最大值求和得到最终分数。

特点:保留了 query 和 document 的 token 级信息,交互比 bi-encoder 更细粒度;document token 嵌入可以预计算和索引(但索引复杂度高于 bi-encoder);计算量介于 bi-encoder 和 cross-encoder 之间。

Score Type 三维对比表格

维度bi-encodercross-encoderlate-interaction
编码方式独立双编码联合编码独立编码 + 后期交互
交互时机无 / 晚期(向量点积)早期(Attention 中)晚期(MaxSim)
相似度计算余弦相似度 / 点积线性分类器输出MaxSim(逐 token 匹配)
速度快(ms 级,支持 ANN)慢(每对一次前向传播)中等(需计算所有 token 对)
可预计算Document 向量可预计算不可预计算Document token 向量可预计算
精度中等高中高
适用阶段粗排(召回)精排(重排序)粗排 / 精排均可
代表模型BERT, GTE, JinaBERT-for-Ranking, JinaRerankerColBERT, ColPali, ColQwen
vLLM 任务embedclassifytoken_embed
典型延迟5ms / 百万文档50ms / 100 文档20ms / 千文档

Score API 处理 score/rerank 请求的实现

vLLM 的 LLM.score(...) API 根据模型的 Score Type 自动选择评分机制:

  • 对于 bi-encoder 模型:分别编码两个输入,计算余弦相似度。
  • 对于 cross-encoder 模型:将两个输入拼接后输入模型,取分类输出。
  • 对于 late-interaction 模型:分别编码后计算 MaxSim。
← 上一模块
vLLM 核心创新
下一模块 →
分布式推理与部署