Prhub

#23838 [Speculative Decoding] Validate vocabulary compatibility in STANDALONE mode

原始 PR 作者 jmamou 合并时间 2026-06-30 05:45 文件变更 1 提交数 16 评论 25 代码增减 +37 / -0

执行摘要

新增 STANDALONE 模式词汇表兼容性验证

STANDALONE speculative decoding requires the draft model to share the same vocabulary as the target model. If the vocabularies differ, draft token IDs map to different strings in the target vocabulary, making the speculative decoding lossy. Previously, SGLang silently accepted any draft model with no error or warning. Concrete example with google/gemma-2-9b (target) and bigcode/tiny_starcoder_py (draft) shows corrupted output.

该 PR 值得快速合并,因为它修复了一个静默输出损坏的严重 bug(虽然仅在特定配置下出现),且影响面小、风险低。review 中关于性能与兼容性的讨论具有通用参考价值。

讨论亮点
  1. get_vocab 性能与兼容性:gemini-code-assist 担心完整词汇表比较对 256k+ tokenizer 的开销,以及 AttributeError 风险。作者测量了 get_vocab 两次 + dict 比较约 314ms,认为在分钟级启动过程中可忽略;通过添加 hasattr 防护解决兼容问题。
  2. TLI 引用移除:kpham-sgl 要求将所有 TLI(Token-Level Interleaved)相关的文档和错误消息引用移至 PR #22883,以保持两个 PR 关注点分离。作者完全移除了文档改动和错误消息中的 TLI 提示。
  3. 方法设计:kpham-sgl 指出 StandaloneWorker 自身拥有 draft 端信息,应使用实例方法而非静态方法。作者将 _validate_vocab_compatibility 改为实例方法,从 self 读取 draft 属性。

实现拆解

  1. 在 StandaloneWorkerV2.init 中,在创建 draft worker 之后、初始化其余张量之前,插入对 _validate_vocab_compatibility 的调用,传入 target 的 vocab_size 和 tokenizer。
  2. 新增实例方法 _validate_vocab_compatibility:
    • 从 self._draft_worker 获取 draft 的 vocab_size 和 tokenizer。
    • 首先比较 target_vocab_size 与 draft_vocab_size,如果不一致则抛出 ValueError,并提示使用兼容的 draft 模型或支持异构词汇表的推测算法。
    • 接着检查两个 tokenizer 是否都非 None 且都有 get_vocab 方法(通过 hasattr 防护),若存在且字典不同则抛出 ValueError。
  3. 方法仅被调用一次,不修改推测解码的核心逻辑,也不影响前向传播或采样。
  4. 配套的单元测试(test_standalone_vocab_check.py)最初被添加以在 CPU 上验证检查逻辑,但由于测试框架变更和后续依赖清理,最终在合并前被移除。
文件 模块 状态 重要度
python/sglang/srt/speculative/standalone_worker_v2.py 推测解码 modified 7.03

关键符号

_validate_vocab_compatibility

关键源码片段

python/sglang/srt/speculative/standalone_worker_v2.py core-logic

唯一修改的文件;新增 _validate_vocab_compatibility 方法并在 __init__ 中调用。

# python/sglang/srt/speculative/standalone_worker_v2.py
def _validate_vocab_compatibility(
    self,
    target_vocab_size: int,
    target_tokenizer,
) -> None:
    """检查 draft 和 target 的词汇表是否兼容,不兼容则抛出 ValueError"""
    # 从 self 读取 draft 端的信息
    draft_vocab_size = self._draft_worker.draft_runner.model_config.vocab_size
    draft_tokenizer = self._draft_worker.draft_worker.tokenizer
​
    # 首先比较大小,这是最快且最常见的失败路径
    if target_vocab_size != draft_vocab_size:
        raise ValueError(
            f"STANDALONE speculative decoding requires the draft model to share the "
            f"same vocabulary as the target model, but got "
            f"target vocab_size={target_vocab_size} and "
            f"draft vocab_size={draft_vocab_size}. "
            f"Use a draft model with a matching vocabulary, or a speculative "
            f"algorithm that supports heterogeneous vocabularies."
        )
​
    # 当两个 tokenizer 都有 get_vocab 方法时才进行完整映射比较
    # 使用 hasattr 防护以兼容不支持该方法的 tokenizer(如 TiktokenTokenizer)
    if (
        target_tokenizer is not None
        and draft_tokenizer is not None
        and hasattr(target_tokenizer, "get_vocab")
        and hasattr(draft_tokenizer, "get_vocab")
        and target_tokenizer.get_vocab() != draft_tokenizer.get_vocab()
    ):
        raise ValueError(
            "STANDALONE speculative decoding requires the draft model to share the "
            "same vocabulary as the target model, but the two tokenizers have "
            "different token-to-id mappings even though their vocab sizes match. "
            "Use a draft model with a matching vocabulary, or a speculative "
            "algorithm that supports heterogeneous vocabularies."
        )

评论区精华

get_vocab 性能与兼容性 性能

gemini-code-assist 提出完整词汇表比较对 256k+ tokenizer 开销大,且 get_vocab 可能在某些 tokenizer 上不存在。

结论:作者测量了每次调用约 314ms(在分钟级启动过程中可忽略),并添加了 hasattr 防护避免 AttributeError。 · 已解决

移除文档中的 TLI 引用 documentation

kpham-sgl 要求将文档改动移到 TLI PR #22883,以保持两个 PR 独立。

结论:作者移除了文档改动和错误消息中的 TLI 引用。 · 已解决

方法设计:staticmethod 改为 instance method 设计

kpham-sgl 指出 StandaloneWorker 拥有 draft_tokenizer 和 draft_vocab_size,不应是静态方法。

结论:作者改为实例方法,从 self 读取 draft 端信息。 · 已解决

风险与影响

低风险。变更仅限于服务启动时的词汇表验证,不修改任何推理路径。潜在风险包括:

  • 如果 tokenizer 不支持 get_vocab 且未用 hasattr 防护,会抛出 AttributeError(已修复)。
  • 性能影响极小(约 314ms),在模型加载和 CUDA 图捕获的背景下可忽略。
  • 行为变更:原先能启动并产生损坏输出的不兼容配置现在会立即报错,但这是预期行为,符合安全改进目标。

对用户:使用 STANDALONE 模式且 draft/target 词汇表不匹配的用户会立即看到 ValueError 启动失败,而不是等待推理输出损坏。匹配的用户无感知。对系统:无运行时性能影响。对团队:该 PR 与 TLI 特性(#22883)解耦,后续 TLI 可以实现异构词汇表支持。

启动验证 无运行时影响 可能误报

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论