Prhub

#45876 [Rust Frontend] Validate tokenized bad_words vocabulary range

原始 PR 作者 reidliu41 合并时间 2026-06-18 10:28 文件变更 2 提交数 1 评论 2 代码增减 +61 / -0

执行摘要

Rust 前端增加 bad_words 词表越界校验

Rust 降级路径 tokenizes bad_words 得到 _bad_words_token_ids,但此前未校验这些 token ID 是否在词表范围内。Python 侧在 SamplingParams.update_from_tokenizer() 中会校验并抛出 bad_words 校验错误。没有此检查时,越界 bad words token ID 会传递到 sampler/logits masking 路径,引发较模糊的 engine 侧错误。

该 PR 是低风险的安全加固,值得合并。虽 reviewer 指出当前无已知触发路径,但仍提升了防御深度和与 Python 侧的一致性。

讨论亮点

Reviewer BugenZhao 指出当前 bad_words_token_ids 的唯一来源是 tokenize_bad_words,若存在越界 token 说明 tokenizer 本身有 bug,但校验仍有助于防御。作者接受了此观点,校验作为额外安全网。

实现拆解

  1. token_ids.rsvalidate_vocab_range() 函数中新增 bad_words 校验分支:添加对 params.bad_words_token_idsvalidate_param 调用,参数名设为 "bad_words",校验边界为 limits.tokenizer_vocab_size,使用 flatten() 展开嵌套的 token IDs 迭代器。
  2. lower.rs 的测试模块中新增测试桩 FixedTokenizer:为支持非空 token 编码结果,新增实现 Tokenizer traitFixedTokenizer,其 encode 返回预设的 token ID 列表,以便构造 out-of-vocab 场景。
  3. 新增测试用例 lower_sampling_params_rejects_out_of_vocab_bad_words:使用 FixedTokenizer 将字符串 "blocked" tokenize 成 [1999, 2000],其中 2000 超出 tokenizer_vocab_size=2000(半开区间,合法为 0..2000),断言返回 Error::OutOfVocab 且参数名、token IDs、vocab_size 正确。
文件 模块 状态 重要度
rust/src/text/src/lower/token_ids.rs 校验层 modified 5.96
rust/src/text/src/lower.rs 测试与桩 modified 7.85

关键符号

validate_vocab_range lower_sampling_params_rejects_out_of_vocab_bad_words

关键源码片段

rust/src/text/src/lower/token_ids.rs core-logic

核心校验逻辑所在,新增 bad_words_token_ids 边界检查分支。

// rust/src/text/src/lower/token_ids.rspub(crate) fn validate_vocab_range(
    params: &SamplingParams,
    limits: &SamplingLimits,
) -> Result<(), OutOfVocabError> {
    // 已有校验: stop_token_ids、allowed_token_ids、logit_bias、logprob_token_ids
    // ...    // 新增:校验 bad_words_token_ids 中每个 token ID 是否 < tokenizer_vocab_size
    if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() {
        validate_param(
            "bad_words", // 参数名,用于错误消息
            bad_words_token_ids.iter().flatten().copied(), // 展开嵌套 Vec<Vec<u32>> 为平坦迭代器
            limits.tokenizer_vocab_size, // 使用 tokenizer 词表大小作为边界
        )?;
    }    Ok(())
}
rust/src/text/src/lower.rs test-coverage

测试 StubTokenizer 扩展为 FixedTokenizer,并新增坏词越界测试用例。

// rust/src/text/src/lower.rs 测试模块// 新增测试桩:返回固定 token ID 列表的 tokenizer
struct FixedTokenizer {
    token_ids: Vec<u32>,
}impl Tokenizer for FixedTokenizer {
    fn encode(&self, _text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
        Ok(self.token_ids.clone()) // 直接返回预设的 token IDs,模拟 tokenize 结果
    }
    fn decode(&self, _token_ids: &[u32], _skip_special_tokens: bool) -> vllm_tokenizer::Result<String> {
        Ok(String::new())
    }
    fn token_to_id(&self, _token: &str) -> Option<u32> {
        None
    }
}#[test]
fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
    // 模拟将 "blocked" tokenize 成 [1999, 2000],其中 2000 超出 tokenizer_vocab_size=2000(合法范围 : 0..2000)
    let tokenizer = FixedTokenizer {
        token_ids: vec![1999, 2000],
    };
    let error = lower_sampling_params(
        SamplingParams {
            bad_words: Some(vec!["blocked".to_string()]),
            ..Default::default()
        },
        SamplingHints::default(),
        sample_sampling_limits(), // tokenizer_vocab_size = 2000
        3,
        &tokenizer,
    )
    .unwrap_err();    // 断言返回 OutOfVocab 错误,参数名 "bad_words",越界 token ID 为 [2000]
    assert!(matches!(
        error,
        Error::OutOfVocab(OutOfVocabError {
            parameter: "bad_words",
            token_ids,
            vocab_size: 2000,
        }) if token_ids == vec![2000]
    ));
}

评论区精华

校验必要性评估 设计

BugenZhao 指出当前 bad_words_token_ids 的唯一来源是 tokenize_bad_words,若存在越界 token 说明 tokenizer 本身有 bug,但校验仍有助于防御。

结论:校验作为防御性编程保留。 · 已解决

风险与影响

变更范围小(2 个文件,61 行新增),仅增加校验逻辑,无回归风险。引入的 FixedTokenizer 仅用于测试,不影响生产代码。

用户:越界的 bad words token IDs 将提前以清晰错误(OutOfVocabError)被拒绝,而非引擎侧模糊失败。系统:无性能影响,校验仅多一次迭代遍历。团队:与 Python 侧行为对齐,降低跨语言行为差异排查成本。

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论