Embedding 与 Reranking 路径#

3.1 到 3.4 讲的都是生成式请求：模型一步步生成 token，直到 stop 或 length 触发。但 SGLang 还支持另外两类完全不同的执行路径：Embedding（把文本压缩成向量）和 Reranking（对查询-文档对打分）。这两类路径没有生成循环，不需要 KV cache，执行逻辑和返回对象都和生成路径有本质差异。

这一节回答三件事：

1. Embedding 和 Reranking 请求走的是什么路径，和生成请求有哪些分叉；
2. ModelRunner 在这两类请求里实际执行的是什么；
3. 输出对象是什么形态，怎样回到 API 表面。

生成 vs Embedding：两种完全不同的任务模式#

先把两种模式的差异压成一张对比表：

最重要的一点是：Embedding 请求只有 prefill，没有 decode 循环。模型跑完一次前向，从 last_hidden_state 或 pooler_output 里取出向量，直接回包。因为没有 decode，KV cache 在这条路径上几乎没有价值。

Embedding 请求的进入路径#

POST /v1/embeddings
    ↓
OpenAIServingEmbedding.create_embedding(...)
    ↓
EmbeddingReqInput(text, ...) — 无 sampling_params，无 max_new_tokens
    ↓
TokenizerManager.embedding_request(...)
    ↓
TokenizedEmbeddingReqInput — tokenize 后的 embedding 输入
    ↓
Scheduler.handle_embedding_request(...)
    ↓
EmbeddingReq — 注意：不是 Req，而是专门的 embedding 请求对象

和生成请求相比，有几处重要差异：

EmbeddingReqInput 没有生成相关字段

@dataclasses.dataclass
class EmbeddingReqInput:
    text: Optional[str] = None
    input_ids: Optional[List[int]] = None
    # 注意：没有 sampling_params, max_new_tokens, stream, stop 等字段
    is_single: bool = True
    normalize: bool = False  # 是否对输出向量做 L2 归一化
    return_sparse: bool = False  # sparse embedding

Scheduler 对 EmbeddingReq 的处理更简单

生成请求需要进 waiting queue 等待 batch 组装。EmbeddingReq 的调度更直接：

• 没有 decode 阶段，不需要跨轮次维护状态；
• 每次 forward 就是完整的请求生命周期；
• 不需要 kv_allocated_len、kv_committed_len 这类 KV 追踪字段。

ModelRunner 怎样执行 Embedding#

当 ForwardBatch.forward_mode == ForwardMode.EMBED 时，ModelRunner.forward() 走的是另一条分支：

if batch.forward_mode == ForwardMode.EMBED:
    # 不进行 token 采样，而是取 hidden states
    hidden_states = self.model.embed(batch)  # [batch, seq_len, hidden_size]
    
    # Pooling: 取最后一个 token 的 hidden state（或做 mean pooling）
    embeddings = pool_hidden_states(hidden_states, batch)  # [batch, hidden_size]
    
    return BatchEmbeddingOutput(embeddings=embeddings)

Pooling 策略：不同的 embedding 模型有不同的 pooling 策略：

• Last token pooling（最常见于 decoder-only 模型如 LLM-Embedder）：取序列最后一个 token 的 hidden state
• Mean pooling（常见于 BERT-style 模型）：对所有 token 的 hidden state 做平均
• CLS token pooling（常见于 encoder-only 模型）：取 [CLS] token（位置 0）的 hidden state

SGLang 通过模型配置里的 pooling_type 字段决定使用哪种策略。

向量维度：embedding 向量的维度等于模型的 hidden_size：

• Llama-3.2 1B：2048 维
• BGE-M3：1024 维
• text-embedding-ada-002 (OpenAI)：1536 维

如果请求里 normalize=True，还会在返回前对每个向量做 L2 归一化：v = v / ||v||₂，这让向量可以直接用点积计算余弦相似度。

Embedding 的输出路径#

BatchEmbeddingOutput 的结构比 BatchStrOutput 简单很多：

@dataclasses.dataclass
class BatchEmbeddingOutput:
    embeddings: List[List[float]]  # [batch_size, hidden_size]
    # 注意：没有 finished_reasons, output_strs 等生成相关字段

因为 embedding 请求没有 decode 循环，也没有 streaming，返回链也比生成请求短得多：

ModelRunner → BatchEmbeddingOutput
    ↓（不经过 DetokenizerManager）
TokenizerManager._handle_batch_output_embedding(...)
    ↓
OpenAI 格式的 embedding response

{
  "object": "list",
  "data": [
    {
      "object": "embedding",
      "embedding": [0.0023, -0.0091, 0.0123, ...],  // 1024 或 2048 个浮点数
      "index": 0
    }
  ],
  "model": "bge-m3",
  "usage": {"prompt_tokens": 15, "total_tokens": 15}  // 没有 completion_tokens
}

Reranking（交叉编码器打分）#

Reranking 的目标是给一个（query, document）对打一个相关性分数，用于搜索结果排序。

API 入口：POST /v1/classify（SGLang 扩展接口）

执行方式：把 query 和 document 拼接后（通常用特殊分隔符如 [SEP]）输入交叉编码器，取 [CLS] 位置的 logit 作为相关性分数。

# 输入拼接格式（以 BGE-Reranker 为例）
query = "What is machine learning?"
document = "Machine learning is a branch of AI..."
input_text = f"[CLS] {query} [SEP] {document} [SEP]"

# 模型前向后取 classification head 的输出
score = model.classify(input_ids)  # 返回 scalar 或 [yes, no] 的 logits

输出：一个浮点数分数（经 sigmoid 变换到 [0, 1] 范围）。

Reranking 和 Embedding 共用很多代码路径（都是一次 forward，都不需要 decode），区别仅在于输出头不同：Embedding 取 hidden state，Reranking 取 classification logit。

encoder-only 模型的特殊支持#

SGLang 通过 --is-embedding-model 标志支持 encoder-only 模型（如 BERT、BGE-M3、E5 等）。这类模型在架构上只有 encoder stack，没有 causal mask，没有 autoregressive 能力。

启动时的差异：

# 普通生成模型
python -m sglang.launch_server --model-path meta-llama/Llama-3-8B --port 30000

# Embedding 模型（encoder-only）
python -m sglang.launch_server --model-path BAAI/bge-m3 --is-embedding-model --port 30001

# 生成模型用作 embedding（decoder-only embedding）
python -m sglang.launch_server --model-path nvidia/NV-Embed-v2 --port 30002
# NV-Embed-v2 是基于 Mistral 的 decoder-only embedding 模型

调试 Embedding 路径时先看哪里#

现象：embedding 向量全为 0 或维度不对

1. 确认 forward_mode 是否为 EMBED；
2. 检查 pooling_type 配置是否和模型架构匹配（decoder-only 模型不应该用 CLS pooling）；
3. 确认模型是否支持 embedding 输出（某些生成模型没有实现 embed() 方法）。

现象：embedding 请求比预期慢很多

Embedding 的耗时主要来自 prefill，没有 decode 开销。如果慢：

1. 检查输入序列长度——长文本 embedding 的 O(S²) attention 成本是主要开销；
2. 检查是否有不必要的 normalize 或 sparse embedding 计算；
3. Embedding 不受益于 KV cache，但可以受益于 batch size 增大（多个文本一起 embed）。

现象：生成模型返回 embedding 时精度很差

decoder-only 生成模型做 embedding 时，需要专门的 instruction tuning（如 NV-Embed-v2 的 “Represent this…“前缀）。直接用未经 embedding 训练的生成模型会得到质量很差的向量。

小结#

• Embedding 请求只有一次 prefill forward，没有 decode 循环，不需要 KV cache；
• EmbeddingReqInput 没有 sampling_params，EmbeddingReq 比 Req 简单得多；
• ForwardMode.EMBED 让 ModelRunner 返回 hidden states 而不是 logits；
• Pooling 策略（last token / mean / CLS）由模型配置决定；
• Reranking 走同样的路径，只是最后取 classification logit 而不是 hidden state；
• encoder-only 模型通过 --is-embedding-model 标志启用，decoder-only 模型也可以做 embedding。

理解了这条路径，就不会把 embedding 请求失败误归到 KV cache 或 sampling 相关的问题上。