执行摘要
本 PR 在 Rust 服务器消息层中新增 SamplingParams 结构体,完整映射 Python 的采样参数定义、默认值和验证流程,并添加了 stop_regex 验证的 admission cache,避免重复 HIR 翻译。这是 Rust 服务器替换 Python TokenizerManager 的关键步骤。
功能与动机
此 PR 从 #29799 拆分,依赖 #32240,目标是在嵌入的 Rust 服务器中独立处理采样参数,从而将 SamplingParams 的创建、验证和 stop 字符串处理从调度器的关键路径移到入口 FSM 步骤,降低 CPU 开销。
实现拆解
- 字段默认值宏:
sampling.rs 中通过 defaulted! 宏为每个字段生成 default() 和 deserialize() 方法,使得显式 null 自动回退到默认值,与 Python SamplingParams.__post_init__ 一致。
- 核心逻辑移植:实现
normalize 和 verify 方法,按顺序执行默认值注入、类型纠正和范围校验;其中 vocab_size 为 Option<u64>,Review 指出需在 None 时抛出异常以匹配 Python 行为。
- 反序列化适配:
SamplingParamsInput 通过 deserialize_any 手动处理单对象/列表两种格式,保留字段级错误信息;SamplingParams 标记 #[serde(deny_unknown_fields)],未知字段直接返回 400。
- 正则缓存优化:
regex.rs 中新增 ADMISSION_CACHE 全局缓存 (LazyLock + Mutex + HashMap),提供 cached_bound 和 cache_bound;RegexPattern::build 优先查缓存,避免重复 HIR 翻译,对包含多个 stop_regex 的请求可节省毫秒级开销。
- 依赖调整:引入
BTreeMap、fmt、serde 详细特征等,为后续完整性校验做准备。
rust/sglang-server/src/utils/regex.rs
新增 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 })
}
}
评论区精华
MortalHappiness: “在 Python 实现中 vocab_size 是必须参数,范围检查始终运行;而 Rust 实现中 vocab_size 为 Option<u64>,可能为 None 导致绕过检查,建议添加检查。”
rainj-me: “Sure will fix the vocab_size and make sure we check the value during boot.”
风险与影响
- 正确性风险:
vocab_size 为 None 时当前不会抛出异常,可能让无效采样参数流入模型导致推理异常。已确认将修复。
- 性能影响:Admission cache 使用
Mutex,但单线程入口下影响极小;缓存最高节省 18ms 验证时间。
- 兼容性:新增字段需与 Python
sampling_params.py 保持同步,否则 msgspec 会静默丢弃未知字段。
- 测试覆盖:两个新模块均无直接单元测试,建议后续补充。
关联脉络
- 依赖 #32240 (Rust 请求消息层基础 PR)
- 从 #29799 拆分,作为 Rust 服务器功能集的独立拼图
- 与 #32242 (Rust 服务器消息类型) 属于同一功能线,共同构成 Rust 嵌入服务器的消息处理层
参与讨论