Prhub

#28744 [router] Tokenize prompt once at ingress; forward input_ids to the engine (all policies)

原始 PR 作者 Kangyan-Zhou 合并时间 2026-06-20 07:07 文件变更 8 提交数 5 评论 4 代码增减 +1620 / -201

执行摘要

sgl-router 入口处 tokenize 一次并复用,消除引擎侧重复 tokenization,最高降低 41% TTFT

PR body 指出:路由器的 cache-aware 策略会 tokenize prompt 计算哈希,引擎随后又会 tokenize 相同 prompt 用于推理,对于 55k-125k+ token 长上下文,重复 tokenization 带来 100-580ms 的冗余延迟。本 PR 旨在消除该冗余,同时为所有策略统一 tokenization 入口。

本 PR 设计质量高,将 tokenize 解耦并安全转发至引擎,强烈建议精读 input_ids_safe_to_forwardrequest_tokens_for 的实现。特别是安全守卫的谓词设计值得参考——它明确枚举了所有不支持转发的情况,未来需持续维护。对于自定义模板等新兴场景,建议建立类似 PR 流程确保安全扩展。

讨论亮点

Review 评论中 gemini-code-assist[bot] 指出:

  • 如果用户请求中指定自定义 chat_template,引擎会使用该模板渲染 prompt,而 router 已用默认模板 tokenize,直接转发 input_ids 会静默绕过自定义模板,导致推理错误。建议将 chat_template 加入阻断键列表(高优先级)。
  • 建议使用 serde_json::to_value(ids) 替代手动构造 Value::Array,代码更简洁且利用 serde 优化(中优先级)。
    同时建议增加对应的测试用例。
    作者在后续 commit(9cf075d)中采纳了 chat_template 阻断建议。关于 serde_json::to_value 的建议未在 diff 中体现,可能已通过其他方式处理。

实现拆解

  1. 在入口处理函数 chat_completionschat.rs)中,根据模型是否有 chat encoder 或策略是否需要 tokens 决定是否解析请求体为 JSON 并调用 request_tokens_for 生成预计算 tokens;
  2. policies/mod.rs 中新增 RequestTokens 结构体和 request_tokens_for 函数,封装 tokenize 逻辑并标定 engine_equivalent 标志,所有策略通过共享 tokenizer 注册表调用;
  3. 修改 CacheAwareZmqPolicyselect 方法(cache_aware_zmq.rs),优先使用 SelectionContext 中的预计算 tokens 而非自行 tokenize,未预计算时回退;
  4. build_outgoing_bodychat.rs)中根据 input_ids_safe_to_forward 检查决定是否携带 input_ids 字段,安全条件包括不包含 toolsmultimodalchat_templatereasoning 等键;
  5. 新增指标 sgl_router_ingress_tokenize_errors_total 和 WARN 日志(metrics.rs),监控 tokenize 失败;
  6. 添加对 chat_template 等阻断键的处理(最后一个 commit),避免转发不正确的 input_ids;
  7. 测试配套:新增 cache_aware_input_ids.rssticky_input_ids.rsroundrobin_input_ids.rs 三个集成测试文件,覆盖各种场景。
文件 模块 状态 重要度
experimental/sgl-router/src/server/routes/chat.rs 请求处理 modified 9.25
experimental/sgl-router/src/policies/mod.rs 策略基础 modified 8.97
experimental/sgl-router/src/policies/cache_aware_zmq.rs 路由策略 modified 8.86
experimental/sgl-router/src/server/metrics.rs 指标系统 modified 7.93
experimental/sgl-router/tests/proxy/sticky_input_ids.rs 测试 added 8.02

关键符号

request_tokens_for input_ids_safe_to_forward build_outgoing_body needs_request_tokens record_ingress_tokenize_error

关键源码片段

experimental/sgl-router/src/server/routes/chat.rs entrypoint

入口处理函数,新增 ingress tokenization 决策、build_outgoing_body 和 input_ids_safe_to_forward 守卫逻辑,改动量最大(+536/-57)

    // Tokenize once at ingress whenever it can pay off — decoupled from the
    // routing policy, because forwarding `input_ids` is a property of the
    // MODEL (does it have a chat encoder so the router can produce
    // engine-equivalent tokens?), not of how we pick the worker. Two gates:
    //
    // * `has_chat_encoder` → a chat request on this model yields
    // engine-equivalent ids we can forward as `input_ids` so the engine
    // skips re-tokenizing. This enables the offload for EVERY policy —
    // sticky and round-robin included — not just cache-aware.
    // * `needs_request_tokens()` → the cache-aware policy ALSO wants the
    // raw-prompt path tokenized for tree matching even on a model with no
    // chat encoder (`/v1/completions` / `text`), which the first gate
    // alone wouldn't trigger.
    //
    // When neither holds, `parse_probe`'s minimal probe is enough, so we keep
    // avoiding the full `serde_json::Value` allocation over a (up to 1 MiB)
    // body. When parsed, this single value is reused for the routing
    // tokenization and the outgoing-body injection below (and PD bootstrap
    // injection). `parse_probe` already validated the object shape.
    let want_tokens = ctx.tokenizers.has_chat_encoder(&model_str)
        || policy.needs_request_tokens();
    let request_value: Option<serde_json::Value> = if want_tokens {
        Some(serde_json::from_slice(&body)
            .map_err(|_| {
                ApiError::BadRequest("invalid request: body must be a JSON object".into())
            })?)
    } else {
        None
    };    // The ids feed both the routing decision (cache-aware consumes them; other
    // policies ignore them) and — when engine-equivalent — the engine itself,
    // forwarded as `input_ids` so it skips re-tokenizing the same prompt. The
    // ingress owns the tokenize via the shared registry, so the choice of
    // policy never changes whether we tokenize.
    let request_tokens = request_value
        .as_ref()
        .and_then(|v| request_tokens_for(&ctx.tokenizers, &model_id, v));
experimental/sgl-router/src/policies/mod.rs entrypoint

新增 RequestTokens 结构体和 request_tokens_for、tokenize_text 等核心函数,作为全局共享的 tokenization 入口

/// Produce the routing tokens — and whether they are engine-equivalent —
/// from an already-parsed request body, using the shared tokenizer registry.
pub fn request_tokens_for(
    tokenizers: &TokenizerRegistry,
    model_id: &ModelId,
    value: &serde_json::Value,
) -> Option<RequestTokens> {
    // Chat-encoder path: only when the model has a chat encoder and the
    // request has `messages`. The engine-equivalent flag is set so the
    // caller knows these ids can be forwarded as `input_ids`.
    if tokenizers.has_chat_encoder(&model_id.0) {
        if let Some(messages) = value.get("messages").filter(|m| m.is_array()) {
            if let Some(ids) = tokenizers.encode_chat(&model_id.0, messages) {
                return Some(RequestTokens {
                    ids,
                    engine_equivalent: true,
                });
            }
        }
    }
    // Raw-prompt fallback: extract text from `prompt` / `text` field
    // and tokenize it. These ids are NOT engine-equivalent (the engine
    // will still apply its own template).
    let text = extract_prompt_text_from_value(value)?;
    let ids = tokenize_text(tokenizers, model_id, &text)?;
    Some(RequestTokens {
        ids,
        engine_equivalent: false,
    })
}

评论区精华

chat_template 应加入阻断列表以避免自定义模板被静默忽略 正确性

gemini-code-assist[bot] 指出若用户指定自定义 `chat_template`,router 已用默认模板 tokenize,转发 input_ids 会导致引擎使用错误模板渲染。建议加入阻断键列表。

结论:作者在后续 commit 中添加了 `chat_template` 到阻断列表,并调整了对应测试。 · 已解决

使用 serde_json::to_value 优化 input_ids 序列化 style

gemini-code-assist[bot] 建议使用 `serde_json::to_value(ids)` 替代手动构造 `Value::Array`,更简洁且利用 serde 优化。

结论:未在后续 patch 中明确体现,可能已通过其他方式处理或保持原样。 · unresolved

风险与影响

  • tokenization 一致性假设:转发 input_ids 依赖 router chat encoder 产生与引擎相同的 tokens,这是 cache-aware 已依赖的前提,guard 机制确保不转发时会回退到引擎 tokenize,但若 encoder 实现有细微信号差异可能导致静默 mismatch。
  • 安全守卫完整性input_ids_safe_to_forward 谓词需覆盖所有可能改变 prompt 的字段(如 chat_templatecontinue_final_message 等),review 中已补充 chat_template,但未来新字段可能遗漏。建议维护清晰的列表。
  • 核心路径变更chat_completions 是请求入口,新增 JSON 解析和 tokenize 路径在长上下文下有性能收益但增加 CPU 内存开销,已在负载测试中验证。
  • 测试覆盖:虽新增集成测试,但对自定义模板阻断等边界场景的测试可能在 review 后补充,需确认。
  • 用户侧:长上下文请求(>55k tokens)的端到端延迟显著降低(idle 最多 -41%,负载下 -34% 至 -49%),短请求无明显损害;所有路由策略均受益。
  • 运维侧:新增指标 sgl_router_ingress_tokenize_errors_total 可监控 tokenization 健康状态;日志 WARN 帮助定位问题。
  • 开发侧:tokenization 逻辑集中到 policies/mod.rs,策略接口更干净;后续新增策略无需关心 tokenize 细节。
  • 影响范围:限于 sgl-router 模块,不影响引擎或其他组件。
核心路径变更 tokenization 一致性假设 安全守卫未全覆盖 自定义模板阻断已修复

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论