Embedding & Pooling：用 SGLang 跑非生成任务

谁该读这一篇？ 要在 SGLang 上同时服务 chat / embedding / rerank 任务的工程师。 前置阅读： 00-prerequisites.md。 耗时： 25 分钟 学完能： 1. 用 SGLang 跑 BGE / E5 / BGE-M3 / Jina 嵌入模型； 2. 解释 pooling 策略（mean / cls / last）； 3. 区分 embedding 和 generation 在调度 / KV 上的差异； 4. 跑 rerank 任务； 5. 评估是否值得让 chat 和 embedding 共用 server。

1. Embedding 任务是什么

Input text ─▶ Tokenize ─▶ Transformer ─▶ Hidden states ─▶ Pool ─▶ Vector [hidden_dim]

输出不是 token id，是 fixed-dim vector。
用于：相似度计算、retrieval、RAG。
没有 decode 循环，只跑一次 forward。

2. SGLang 启动 embedding server

python -m sglang.launch_server \
    --model-path BAAI/bge-large-en-v1.5 \
    --is-embedding \
    --tp 1 \
    --port 30000

--is-embedding 切换到 embedding 模式。

OpenAI 兼容 API：

from openai import OpenAI
client = OpenAI(base_url="http://localhost:30000/v1", api_key="EMPTY")
resp = client.embeddings.create(
    model="BAAI/bge-large-en-v1.5",
    input=["text 1", "text 2"]
)
print(resp.data[0].embedding[:5])    # [0.012, -0.034, ...]

3. Pooling 策略

不同模型用不同 pooling：

策略	公式	用模型
CLS	取 `[CLS]` token 的 hidden state	BERT 系
Mean	平均所有 token	E5 / BGE
Last	取最后一 token	GPT-style embedding（Llama embedding）
Weighted	attention-weighted	自定义

SGLang 在 --pool 或 modelconfig 里设：

--pool mean
--pool last

模型架构通常硬编码 pooling 方式，自动选对。

4. Embedding 在 SGLang 内部

sequenceDiagram
    participant C as Client
    participant T as TokenizerManager
    participant S as Scheduler
    participant W as ModelWorker
    participant TM as TokenizerManager

    C->>T: POST /v1/embeddings
    T->>S: EmbeddingReqInput
    S->>W: forward(extend mode, output_hidden=True)
    W-->>S: hidden_states
    S->>S: pool → embedding vector
    S-->>T: BatchEmbeddingOutput
    T-->>C: JSON 响应

差异：

不走 detokenizer（没有 token 流）。
不需要 KV cache 累计（一次 forward 结束）。
Scheduler 用 forward_mode=extend，但跑完不进 decode。

5. 与 chat 模型共用 server？

理论上可以同 process 跑 chat + embedding，但 SGLang 当前更推荐：

chat server 一个进程。
embedding server 另一个进程。
客户端 / router 按任务类型路由。

理由：

模型架构差异（embedding 模型一般小、是 encoder-only）。
调度策略不一样（embedding 不需 cache-aware）。
资源配比不同。

如果业务流量小：用同进程的 chat 模型直接 embedding（用 last pooling，[Llama embedding]）也行，质量略低但管控简单。

6. BGE-M3 等高级 embedding

BGE-M3 输出多种 embedding：

Dense（标准向量）
Sparse（稀疏，类似 SPLADE）
ColBERT（每 token 一个 vector）

SGLang 支持 BGE-M3 但要看版本；启动加 --enable-multi-vector。

7. Rerank 任务

输入：query + N 个候选 document，输出：每个 doc 的相关性 score。

# OpenAI 没有 rerank API，SGLang 用 /v1/score 或 native API
resp = requests.post(
    "http://localhost:30000/v1/score",
    json={
        "model": "BAAI/bge-reranker-large",
        "query": "什么是 RadixAttention？",
        "documents": ["doc 1...", "doc 2...", "doc 3..."]
    }
)
print(resp.json()["scores"])

Cross-encoder 内部把 (query, doc) 拼起来做 forward，输出分数。延迟比 dense embedding 高（每对 query-doc 一次 forward）。

8. 性能特征

模型	Latency p50 (H100)	Throughput
BGE-large-en (335M)	5-10 ms / item	5000 + items/s
BGE-M3 (560M)	10-20 ms / item	2000+ items/s
BGE-reranker-large	50-100 ms / pair	200+ pairs/s
Llama-style embedding (7B)	50-100 ms / item	200+ items/s

Embedding throughput >> chat throughput（因为没有 decode）。

9. 批量请求

embedding 天然适合大 batch：

client.embeddings.create(
    model="...",
    input=[t1, t2, t3, ..., t256]      # 一次 256
)

SGLang 内部 batch 一次 forward 跑完，throughput 极高。

10. 模型选型

任务	推荐
通用英文	BGE-large-en-v1.5
通用中文	BGE-large-zh-v1.5
多语言	BGE-M3 / multilingual-e5
长文档	nomic-embed-text-v1.5 (8k context)
Rerank	BGE-reranker-large
极致质量	Voyage / Cohere 闭源

11. 常见问题

现象	排查
启动报 "not embedding model"	模型架构不是 encoder-only；加 `--is-embedding` 强制
Embedding 维度不对	检查 model config；可能加了 projection
同输入输出不一致	数值精度问题（fp16）；用 fp32 验证
慢	batch 太小；增 input list
输出全 0	模型加载失败；查启动日志

12. 监控

sglang:embedding_throughput        # vectors/s
sglang:embedding_latency_ms

健康线：

BGE-large：throughput > 3000/s（H100 batch 32）。
单次 latency p99 < 50 ms。

13. 小结

Embedding 任务一次 forward 出 vector，无 KV cache 累计。
Pooling 策略 cls / mean / last 各对应不同模型。
SGLang 用 --is-embedding 切换模式。
Rerank 是 cross-encoder，延迟更高。
通常 chat 和 embedding 独立部署，按业务路由。

14. 自检

Embedding 为什么不进 detokenizer？

答案

Embedding 任务的输出是 **fixed-dim vector**（如 `[1024]` float），不是 token 流。 Detokenizer 的工作是 token id → 字符串（增量 BPE 解码），对 vector 输出没意义。 Embedding 流程： (1) TM tokenize prompt → token id 列表。 (2) Scheduler 一次 forward 拿 hidden states。 (3) Pooling（cls / mean / last）→ 单个 vector。 (4) Scheduler 直接发 `BatchEmbeddingOutput` 给 TM（**跳过 Detokenizer**）。 (5) TM 返回 JSON 给客户端。架构上 Embedding 走的是 "Scheduler → TM" 直接回流路径，少一跳。

CLS / mean / last pooling 各对应什么类型模型？

答案

- **CLS pooling**：取 `[CLS]` token 的 hidden state。**BERT 系**（encoder-only，bidirectional）：BERT、RoBERTa、早期 sentence-BERT。CLS token 在训练时被设计为 "represent the whole sentence"。 - **Mean pooling**：所有 token hidden state 平均。**E5 / BGE 系**（modern 中英文 embedding 模型）：训练时 in-batch negative + cosine similarity loss，mean 比 CLS 经验上更稳。 - **Last pooling**：取最后一个 token 的 hidden state。**GPT-style / LLM-based embedding**：Llama embedding、Mistral embedding、E5-Mistral。LLM 是 decoder-only causal，只有最后位置看到全文，自然成为"整段摘要"。选错 pooling：embedding 完全错（数值合理但语义不对），cosine similarity 接近随机。 SGLang 按模型自动选；自定义用 `--pool` 覆盖。

BGE-M3 的 dense / sparse / ColBERT 差别？

答案

BGE-M3 一次 forward 输出三种 embedding（multi-vector）： - **Dense**：标准 `[1024]` float 向量。语义相似度（cosine）；快、通用。 - **Sparse**：`Dict[token_id → weight]` 稀疏权重，类似 SPLADE。**词汇匹配**——对关键词检索友好，能精确匹配数字、人名、专有名词；dense 容易遗漏。 - **ColBERT**：每 token 一个 `[128]` 向量，文档侧存 N×128，query 侧 M×128，相似度 = max over token-pair similarity。**细粒度匹配**——长文档精度高但存储 N×128 倍。实战：dense 做粗排（百万级文档秒级返回 top-100），ColBERT 做精排（top-100 精排到 top-10），sparse 作为关键词兜底。三者融合的检索效果比单一 dense 高 20-30%。

Rerank vs dense embedding 在 RAG 里怎么配？

答案

两阶段检索是 RAG 标准做法： - **阶段 1 - Dense embedding 召回**：业务向量库（如 Milvus）按 cosine similarity 召回 top-100 文档；快（百万级文档毫秒）。 - **阶段 2 - Cross-encoder rerank**：对召回的 100 个 (query, doc) 对跑 cross-encoder（如 BGE-reranker-large），输出相关性 score，按 score 排序取 top-5。差别： - Dense（bi-encoder）：query 和 doc 独立编码，速度快但精度低（特别是细微匹配）。 - Rerank（cross-encoder）：(query, doc) 拼起来一次 forward 算 attention，精度高 10-20%（特别是数字、否定句、refined 语义）。开销：rerank 100 对每对 50-100ms，串行慢；批量送 SGLang `/v1/score` API 一次 forward 处理。配置：dense 取宽（top-100）保 recall，rerank 取窄（top-5）保 precision。

同一 server 同时跑 chat + embedding 的代价？

答案

技术上能做但不推荐： (1) **模型架构不同**：chat 是 decoder（causal mask），embedding 通常是 encoder（bidirectional 或 last pooling）。同 server 只能选一种架构。 (2) **调度策略冲突**：chat 用 cache-aware（前缀复用关键），embedding 单次 forward 没复用价值，cache_aware 无意义甚至成负担。 (3) **batch 形状打架**：embedding 大 batch（128+）单次 forward，chat 多请求小 batch + decode 循环；混 batch 让两边都不优化。 (4) **资源配比难**：chat 关心 KV pool（big），embedding 关心 throughput（不需要 KV），同副本配置妥协。推荐：**分独立副本组**。embedding 用专门小副本（BGE-large 单卡能跑 5000+ token/s），chat 用 70B 大模型副本。Router 按 endpoint 路由（`/v1/embeddings` vs `/v1/chat/completions`）。只有低流量场景（< 50 QPS each）才考虑共用，省 1-2 个副本钱不值得复杂度。

15. 下一步

07-hands-on/02-benchmark-throughput.md — 跑 embedding 基准。
08-production-deployment/ — 生产部署。
源码：srt/managers/tokenizer_manager.py embedding 路径、srt/entrypoints/openai/ embedding API。

上游源码：sglang/python/sglang/srt/ embedding 相关分布在 managers / models 各处。