Prhub

#22203 [Spec][Ngram] Support multiple SAMs with dynamic HTTP API

原始 PR 作者 hnyls2002 合并时间 2026-04-07 09:49 文件变更 14 提交数 29 评论 17 代码增减 +684 / -120

执行摘要

为 Ngram 推测解码添加多 SAM 动态 HTTP API 支持,允许运行时管理外部语料库。

动机源于Ngram推测解码的灵活性需求:启动时通过--speculative-ngram-external-corpus-path加载的单个外部SAM无法满足运行时动态管理语料库的需求。用户希望在不重启服务器的情况下添加或移除语料库,以适配不同场景。此PR是Ngram重构系列(Issue #21052)的一部分,跟随PR #21425。

建议技术管理者和工程师精读此PR,重点关注:

  1. 多SAM存储的设计决策,如何通过std::unordered_map管理语料库生命周期和并发限制(见FIXME)。
  2. 异步加载模式,使用ExternalCorpusManager和后台线程实现非阻塞操作,避免影响调度器事件循环。
  3. 预算分配逻辑在batchMatch中的实现,确保草案生成质量不受多SAM影响,注意整数除法可能导致的余数分配问题。
  4. 向后兼容处理,保持启动参数--speculative-ngram-external-corpus-path工作,并将其整合到新API中。
讨论亮点

Review中核心讨论点:

  • 异常处理bug(正确性):reviewer指出在ngram_corpus.pyload_external_corpus_named中,异常时调用clear_external_corpus()会清除所有语料库,而非仅失败的那个。结论:添加了FIXME注释,但未完全解决;建议使用专用清理方法。
  • 预算分配计算(性能):reviewer发现batchMatchtrie_budget计算可能错误,整数除法余数被分配给trie而非SAM间分配。结论:提交中调整计算,将循环不变预算移出循环,但分配逻辑可能仍需优化。
  • Mutex保护注释(文档):reviewer建议恢复和更新ngram.h中mutex_的注释以明确保护范围。结论:提交中恢复了注释,并更新以包括新成员。
  • 字符串连接性能(性能):reviewer指出ngram_corpus_ffi.cpplist_external_corpora使用循环连接字符串可能低效,建议预分配。结论:提交中改用换行符分隔以避免逗号问题,但性能优化未显式实施。

实现拆解

实现分为多个层次:

  1. C++核心层:在python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp中,将单个sam_指针替换为std::unordered_map<std::string, std::unique_ptr<SuffixAutomaton>> sams_,支持多SAM存储;新增staging_sam_用于后台加载;修改batchMatch方法,将总external_sam_budget平均分配给所有活跃SAM,并合并结果。
  2. FFI接口:在ngram_corpus_ffi.cpp中添加remove_external_corpuslist_external_corpora等方法,暴露给Python层。
  3. Python层:更新ngram_corpus.py中的NgramCorpus类,添加load_external_corpus_namedremove_external_corpuslist_external_corpora方法。
  4. HTTP服务器:在http_server.py中新增三个端点:POST /add_external_corpus(支持文件路径或文档列表)、POST /remove_external_corpusGET /list_external_corpora,使用可选管理员认证。
  5. 调度器和令牌管理器:引入ExternalCorpusManager(新增文件external_corpus_manager.py)处理异步加载和响应转发;在scheduler.pytokenizer_communicator_mixin.py中集成请求处理逻辑,沿用flush_cache的通信模式。
  6. 测试:扩展test_ngram_corpus.pytest_server_args.py,验证多SAM操作、预算分配和HTTP API正确性。
文件 模块 状态 重要度
python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp jit-kernel modified 9.0
python/sglang/srt/entrypoints/http_server.py http-server modified 8.0
python/sglang/srt/speculative/external_corpus_manager.py speculative added 7.0
python/sglang/srt/speculative/ngram_worker.py speculative modified 6.0
test/registered/unit/spec/test_ngram_corpus.py test modified 5.0

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

Ngram::startExternalCorpusLoad Ngram::batchMatch add_external_corpus HTTP endpoint ExternalCorpusManager::add NgramCorpus::load_external_corpus_named

评论区精华

异常处理清除所有语料库 正确性

Reviewer 指出在 `ngram_corpus.py` 的 `load_external_corpus_named` 中,异常时调用 `clear_external_corpus()` 会错误清除所有语料库,而非仅失败的那个。

结论:添加了 FIXME 注释,但未完全解决;建议使用专用清理方法或重新开始加载。 · partially_resolved

SAM 预算分配计算 性能

Reviewer 发现 `batchMatch` 中 `trie_budget` 计算可能错误,因为整数除法导致余数被分配给 trie,而非在 SAM 间分配。

结论:提交中调整了计算,将循环不变预算移出循环,但分配逻辑可能仍需优化。 · 已解决

Mutex 保护注释 documentation

Reviewer 建议恢复和更新 `ngram.h` 中 mutex_ 的注释以明确保护范围,包括新成员 `sams_` 和 `staging_sam_`。

结论:提交中恢复了注释,并更新以反映保护范围。 · 已解决

字符串连接性能 性能

Reviewer 指出 `ngram_corpus_ffi.cpp` 中 `list_external_corpora` 使用循环连接字符串可能低效,建议预分配容量。

结论:提交中改用换行符分隔以避免逗号问题,但性能优化未显式实施。 · partially_resolved

风险与影响

技术风险包括:

  1. 并发加载限制:C++代码中FIXME指出staging_sam_为单staging slot,不支持并发语料库加载,可能导致竞争或阻塞用户请求。
  2. 预算分配不均:SAM预算平均分配可能未考虑语料库大小或质量,影响草案生成效果和推测解码性能。
  3. HTTP API安全:新增端点仅受可选管理员认证(AuthLevel.ADMIN_OPTIONAL),若服务器配置不当可能暴露未授权操作。
  4. 回归风险:修改核心匹配逻辑batchMatch,可能影响Ngram推测解码的正确性和性能,尤其是预算分配和结果合并部分。
  5. 测试覆盖不足:尽管有单元测试,但HTTP端点和异步加载的集成测试可能不足,难以覆盖高并发或错误场景。

影响范围:

  • 用户:提供动态管理外部语料库的能力,增强Ngram推测解码的灵活性和适用场景,用户无需重启服务器即可适配不同语料库。
  • 系统:引入非阻塞后台加载减少对推理请求的干扰;多SAM存储可能增加内存开销,但通过预算分配控制总token数;HTTP API新增端点增加服务器维护负担。
  • 团队:新增API和模块(如ExternalCorpusManager)增加代码复杂度,但设计借鉴现有模式(如flush_cache),便于团队理解和维护。
并发加载限制 预算分配不均 HTTP API 安全 核心路径变更

关联 Issue

#21052 [Roadmap] Further Ngram Speculative Decoding Support

完整报告

执行摘要

此PR为Ngram推测解码引入了多SAM(后缀自动机)动态HTTP API支持,允许用户在运行时通过新增的HTTP端点(添加、移除、列出)管理外部语料库,无需重启服务器。实现包括C++层多SAM存储、非阻塞后台加载和预算平均分配,保持向后兼容。这是一个有意义的改进,增强了系统灵活性,但需关注并发限制和预算分配风险。

功能与动机

动机源于Ngram推测解码的灵活性不足问题:之前仅支持通过启动参数--speculative-ngram-external-corpus-path静态加载单个外部SAM,用户无法在运行时动态调整语料库。根据PR body,此变更是Ngram重构系列(Issue #21052)的一部分,跟随PR #21425,旨在提升用户体验和系统适应性。关键表述:"Users may want to add/remove corpora at runtime without restarting the server."

实现拆解

实现按模块拆解如下:

模块 关键改动 说明
C++核心层 (ngram.cpp) 替换sam_sams_映射;新增staging_sam_;修改batchMatch分配预算 核心逻辑变更,支持多SAM存储和匹配
FFI接口 (ngram_corpus_ffi.cpp) 添加remove_external_corpuslist_external_corpora方法 暴露C++功能给Python层
Python层 (ngram_corpus.py) 更新NgramCorpus类,支持命名语料库操作 封装底层调用,提供高级API
HTTP服务器 (http_server.py) 新增/add_external_corpus/remove_external_corpus/list_external_corpora端点 提供外部管理接口,支持文件或文档输入
调度器 (scheduler.py) 集成ExternalCorpusManager,添加请求处理逻辑 协调异步加载和请求转发
令牌管理器 (tokenizer_communicator_mixin.py) 新增通信器处理语料库操作 沿用现有模式,确保数据流一致
测试 (test_ngram_corpus.py) 扩展单元测试,覆盖多SAM功能和HTTP API 验证正确性和回归防护

关键代码片段示例(来自ngram.cpp):

void Ngram::batchMatch(...) {
    // 计算预算分配
    const size_t num_sams = sams_.size();
    const size_t total_sam_budget = num_sams > 0 ? std::min(param_.external_sam_budget, total_draft_token_num) : size_t{0};
    const size_t per_sam_budget = num_sams > 0 ? total_sam_budget / num_sams : size_t{0};
    const size_t trie_budget = total_draft_token_num - (per_sam_budget * num_sams);
    // 后续合并各SAM结果
}

评论区精华

Review讨论由gemini-code-assist[bot]发起,重点包括:

  • 异常处理bug:在ngram_corpus.py中,异常时调用clear_external_corpus()会清除所有语料库,reviewer指出这是错误行为,建议专用清理。结论:添加FIXME注释,但未完全解决。

  • 预算分配计算:reviewer发现trie_budget计算可能错误,整数除法余数被分配给trie。提交中已调整,但分配逻辑仍需关注。引用原话:"The difference (total_sam_budget % num_sams) is currently given to the trie, which might not be the intended behavior."

  • Mutex保护注释:reviewer建议恢复注释以明确保护范围。结论:已更新注释,增强可维护性。

  • 字符串连接性能:reviewer指出list_external_corpora中循环连接字符串可能低效。结论:改用换行符分隔,但性能优化未实施。

风险与影响

风险

  1. 并发加载限制:C++代码中FIXME标记单staging slot,不支持并发加载,可能导致竞争或阻塞。
  2. 预算分配不均:SAM预算平均分配可能忽略语料库特性,影响草案生成质量。
  3. HTTP API安全:端点依赖可选认证,配置不当可能引发未授权访问。
  4. 回归风险:核心batchMatch逻辑变更可能影响推测解码正确性和性能。

影响

  • 用户:获得动态管理能力,提升使用灵活性。
  • 系统:非阻塞加载减少干扰,但多SAM增加内存开销;HTTP API扩展服务器功能。
  • 团队:新增模块增加维护负担,但设计模式借鉴现有组件,降低学习成本。

关联脉络

此PR是Ngram推测解码演进的关键一步:

  • 直接关联:PR #21425引入了外部语料库加载基础,此PR在其上扩展为动态多SAM管理。
  • 系列关联:PR #22180、#22199等同属Ngram优化系列,聚焦性能改进和测试,反映团队持续投入推测解码功能增强。
  • 趋势:从静态加载到动态API,显示系统向更灵活、可运维方向演进,支持运行时调整以适应多变场景。

参与讨论