Prhub

#32343 sglang rust server sampling message

原始 PR 作者 rainj-me 合并时间 2026-07-30 05:55 文件变更 2 提交数 1 评论 3 代码增减 +1036 / -27

执行摘要

Rust 服务器新增采样消息结构体与验证

为了在嵌入的 Rust 服务器中替代 Python TokenizerManager,需要将 SamplingParams 的创建和验证移到 Rust 端。本 PR 是实现该目标的步骤之一,分割自大 PR #29799,并基于 #32240 的消息层基础。

本 PR 是 Rust 服务器构建的基石之一,值得学习如何将 Python 参数验证逻辑系统性地移植为 Rust,以及 admission cache 的设计 (纯函数缓存、容量满时清空策略)。建议 Review 过程中特别关注 vocab_size 范围的正确性。

讨论亮点

Review 中 MortalHappiness 指出:在 Python 实现中 vocab_size 是必须参数,其范围检查始终运行;而 Rust 实现中 vocab_sizeOption<u64>,可能为 None 导致绕过范围检查,建议添加检查。作者 rainj-me 确认将修复。该讨论保证了与 Python 行为一致。

实现拆解

  1. 定义字段默认值宏和常量:在 sampling.rs 中通过 defaulted! 宏批量生成每个字段的 defaultdeserialize 方法,处理 null 回退默认值;同时定义 max_new_tokens_default 返回 Some(128),与 Python 的 max_new_tokens: Optional[int] = 128 一致。
  2. 实现 normalizeverifynormalize 方法按顺序执行默认值注入 (配置超时、stop 字符串处理等) 和数据类型纠正;verify 方法校验所有字段范围。其中 vocab_size 参数在 Python 中为必须参数,Rust 实现中需确保其为 Some 时执行范围检查,None 时需抛出异常。
  3. 自定义反序列化入口SamplingParamsInput 通过 deserialize_any 手动分发单对象或列表,保留内部错误信息。SamplingParams 结构体通过 #[serde(deny_unknown_fields)] 配合 max_new_tokens 的特殊默认函数,实现未知字段 400 响应。
  4. 添加 admission cache:在 regex.rs 中新增 ADMISSION_CACHE 全局缓存 (LazyLock+Mutex+HashMap),提供 cached_boundcache_bound 函数;RegexPattern::build 先查缓存,命中直接返回,否则验证后缓存结果。缓存容量 512 (与 CPython 的 re._MAXCACHE 一致),满时全量清空。
  5. 调整导入和依赖sampling.rs 引入 BTreeMapfmt、serde 详细特征和 OneOrManyRegexPatternregex.rs 引入 HashMapLazyLockMutex
文件 模块 状态 重要度
rust/sglang-server/src/message/sampling.rs 采样配置 modified 8.84
rust/sglang-server/src/utils/regex.rs 正则验证 modified 8.1

关键符号

normalize verify max_new_tokens_default cached_bound cache_bound

关键源码片段

rust/sglang-server/src/utils/regex.rs core-logic

新增 admission cache,缓存已验证的 stop regex 模式,避免重复 HIR 翻译,优化性能

// Admission cache for stop regex patterns.
// Translating a 256-byte \W-heavy pattern measures 574 us; caching avoids
// redoing that on every request.
use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};const ADMISSION_CACHE_CAP: usize = 512;static ADMISSION_CACHE: LazyLock<Mutex<HashMap<Box<str>, usize>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));/// Look up a cached bound for `pattern`.  Returns `None` if not cached.
fn cached_bound(pattern: &str) -> Option<usize> {
    ADMISSION_CACHE.lock().ok()
        .and_then(|c| c.get(pattern).copied())
}/// Store `max_len` for `pattern` in the cache.
// Clears the entire cache when it reaches capacity (like CPython's re._MAXCACHE).
fn cache_bound(pattern: &str, max_len: usize) {
    let Ok(mut c) = ADMISSION_CACHE.lock() else { return; };
    if c.len() >= ADMISSION_CACHE_CAP {
        c.clear();
    }
    c.insert(pattern.into(), max_len);
}impl<'a> RegexPattern<'a> {
    /// Build a `RegexPattern`, checking the cache first.
    fn build(pattern: &'a str) -> Result<Self, Error> {
        // 相同的 pattern 文本必然得出相同的判决和 bound,
        // 因此缓存命中可以避免一次 parse + translate
        if let Some(max_len) = cached_bound(pattern) {
            return Ok(Self { pattern, max_len });
        }
        let ast = validate(pattern)?;
        // … translate HIR and compute bound …
        let max_len = regex_max_seq_length(&ast);
        cache_bound(pattern, max_len);
        Ok(Self { pattern, max_len })
    }
}

评论区精华

vocab_size 检查缺失 正确性

MortalHappiness 指出在 Python 中 `vocab_size` 是必须参数,范围检查始终运行;而 Rust 实现中 `vocab_size` 为 `Option<u64>`,`None` 时会绕过检查,建议修复。

结论:作者 rainj-me 同意添加检查,在后续提交中修复。 · 已解决

风险与影响

vocab_sizeNone 时当前实现不会抛出异常,可能让非法采样参数流入模型,导致推理异常。需确保在 normalizeverify 中处理 None 情况。Admission cache 使用全局 Mutex,在极端并发下可能有轻微锁竞争,但单线程入口设计使其影响极小。缺少单元测试覆盖这两个新模块。

对正在开发中的 Rust 服务器是关键的模块补齐;对其他 Python 业务无影响。开发者需注意保持 sampling.rssampling_params.py 的字段同步,避免 msgspec 静默丢弃未知字段。

vocab_size 验证遗漏 缺少测试覆盖

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论