@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5ed3faa43d

[Rust Frontend] Add ordinary-text tokenizer encoding (#49992)

Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jul 28, 2026 at 15:46 UTC 5ed3faa43ddf075f24482396b634edc33e047a40
10 files changed +448 -2
rust/src/chat/src/renderer/inkling/tests.rs
+4
@@ -31,6 +31,10 @@ impl Tokenizer for FixtureTokenizer {
31 Ok(text.bytes().map(u32::from).collect())
32 }
33
34 + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
35 + self.encode(text, false)
36 + }
37 +
38 fn decode(
39 &self,
40 token_ids: &[u32],
rust/src/parser/benches/utils/adapter.rs
+4
@@ -19,6 +19,10 @@ impl Tokenizer for BenchTokenizer {
19 Ok(text.chars().map(|_| u32::MAX).collect())
20 }
21
22 + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
23 + self.encode(text, false)
24 + }
25 +
26 fn decode(
27 &self,
28 token_ids: &[u32],
rust/src/parser/src/unified/inkling.rs
+8
@@ -414,6 +414,10 @@ mod tests {
414 Ok(text.chars().map(u32::from).collect())
415 }
416
417 + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
418 + self.encode(text, false)
419 + }
420 +
421 fn decode(
422 &self,
423 token_ids: &[u32],
@@ -733,6 +737,10 @@ mod tests {
737 Ok(vec![])
738 }
739
740 + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
741 + self.encode(text, false)
742 + }
743 +
744 fn decode(
745 &self,
746 _token_ids: &[u32],
rust/src/text/src/backend/hf/mod.rs
+4
@@ -177,6 +177,10 @@ mod tests {
177 Ok(vec![])
178 }
179
180 + fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
181 + self.encode(text, false)
182 + }
183 +
184 fn decode(
185 &self,
186 _token_ids: &[u32],
rust/src/tokenizer/src/hf.rs
+288 -2
@@ -1,13 +1,20 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +use std::borrow::Cow;
5 use std::path::Path;
5 -use std::sync::Arc;
6 +use std::sync::{Arc, LazyLock};
7
8 use fastokens::Tokenizer as FastokensTokenizer;
9 use fastokens::decoders::Decoder as FastokensDecoder;
10 +use fastokens::pre_tokenized::{
11 + PreTokenizedString as FastokensPreTokenizedString, Split as FastokensSplit,
12 +};
13 +use fastokens::{PreTokenizer as FastokensPreTokenizer, Split as FastokensSplitPreTokenizer};
14 use thiserror_ext::AsReport as _;
10 -use tokenizers::Tokenizer as HfTokenizer;
15 +use tokenizers::{
16 + AddedVocabulary, Model as _, OffsetType, PreTokenizer as _, Tokenizer as HfTokenizer,
17 +};
18 use tracing::{info, warn};
19
20 use crate::byte_level_decode::decode_byte_level;
@@ -16,6 +23,8 @@ use crate::{Result, Tokenizer};
23
24 mod added_tokens;
25
26 +static EMPTY_HF_ADDED_VOCABULARY: LazyLock<AddedVocabulary> = LazyLock::new(AddedVocabulary::new);
27 +
28 enum Backend {
29 Hf(Box<HfTokenizer>),
30 Fastokens(Box<FastokensTokenizer>),
@@ -53,6 +62,85 @@ fn decode_fastokens_byte_level(
62 Ok(decode_byte_level(tokens))
63 }
64
65 +fn encode_hf_ordinary(tokenizer: &HfTokenizer, text: &str) -> tokenizers::Result<Vec<u32>> {
66 + let mut pretokenized =
67 + EMPTY_HF_ADDED_VOCABULARY.extract_and_normalize(tokenizer.get_normalizer(), text);
68 +
69 + if let Some(pre_tokenizer) = tokenizer.get_pre_tokenizer() {
70 + pre_tokenizer.pre_tokenize(&mut pretokenized)?;
71 + }
72 + pretokenized.tokenize(|normalized| tokenizer.get_model().tokenize(normalized.get()))?;
73 + let encoding = pretokenized.into_encoding(None, 0, OffsetType::Byte)?;
74 + let encoding = tokenizer.post_process(encoding, None, false)?;
75 + Ok(encoding.get_ids().to_vec())
76 +}
77 +
78 +fn fastokens_fused_split(tokenizer: &FastokensTokenizer) -> Option<&FastokensSplitPreTokenizer> {
79 + // Keep this predicate aligned with fastokens::Tokenizer::detect_fused_byte_level.
80 + let FastokensPreTokenizer::Sequence(steps) = tokenizer.pre_tokenizer()? else {
81 + return None;
82 + };
83 + let [
84 + FastokensPreTokenizer::Split(split),
85 + FastokensPreTokenizer::ByteLevel(byte_level),
86 + ] = steps.as_slice()
87 + else {
88 + return None;
89 + };
90 + byte_level.is_bulk_only().then_some(split)
91 +}
92 +
93 +fn fastokens_pre_tokenized_ordinary(
94 + tokenizer: &FastokensTokenizer,
95 + text: &str,
96 +) -> FastokensPreTokenizedString {
97 + // This is fastokens::Tokenizer::build_pre_tokenized with added_tokens = None.
98 + let normalized = tokenizer
99 + .normalizer()
100 + .map_or(Cow::Borrowed(text), |normalizer| normalizer.normalize(text));
101 + match normalized {
102 + Cow::Borrowed(_) => FastokensPreTokenizedString::from_text(text),
103 + Cow::Owned(text) => {
104 + let len = text.len();
105 + FastokensPreTokenizedString::new(
106 + text,
107 + vec![FastokensSplit {
108 + range: 0..len,
109 + token_id: None,
110 + }],
111 + )
112 + }
113 + }
114 +}
115 +
116 +fn encode_fastokens_ordinary(
117 + tokenizer: &FastokensTokenizer,
118 + text: &str,
119 +) -> std::result::Result<Vec<u32>, fastokens::Error> {
120 + if text.is_empty() {
121 + return Ok(Vec::new());
122 + }
123 +
124 + let mut pretokenized = fastokens_pre_tokenized_ordinary(tokenizer, text);
125 + let ids = if let Some(split) = fastokens_fused_split(tokenizer) {
126 + split.pre_tokenize(&mut pretokenized)?;
127 + pretokenized
128 + .tokenize_batched(|buffer, splits, output| {
129 + tokenizer.model().tokenize_batch_fused(buffer, splits, output)
130 + })
131 + .map_err(fastokens::Error::Model)?
132 + } else {
133 + if let Some(pre_tokenizer) = tokenizer.pre_tokenizer() {
134 + pre_tokenizer.pre_tokenize(&mut pretokenized)?;
135 + }
136 + pretokenized
137 + .tokenize(|text, output| tokenizer.model().tokenize_into(text, output))
138 + .map_err(fastokens::Error::Model)?
139 + };
140 +
141 + Ok(tokenizer.post_process(ids, false))
142 +}
143 +
144 /// Tokenizer from `tokenizer.json` in HuggingFace format.
145 ///
146 /// This tries to load with `fastokens` first for better performance, then falls
@@ -156,6 +244,17 @@ impl Tokenizer for HuggingFaceTokenizer {
244 }
245 }
246
247 + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> {
248 + match &self.backend {
249 + Backend::Hf(tokenizer) => encode_hf_ordinary(tokenizer, text)
250 + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report())),
251 + Backend::Fastokens(tokenizer) | Backend::FastokensByteLevel(tokenizer) => {
252 + encode_fastokens_ordinary(tokenizer, text)
253 + .map_err(|error| tokenizer_error!("encoding failed: {}", error.as_report()))
254 + }
255 + }
256 + }
257 +
258 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
259 match &self.backend {
260 Backend::Hf(t) => t
@@ -200,12 +299,19 @@ impl Tokenizer for HuggingFaceTokenizer {
299
300 #[cfg(test)]
301 mod tests {
302 + use std::path::{Path, PathBuf};
303 +
304 + use serde_json::{Value, json};
305 use tempfile::tempdir;
306 use tokenizers::models::bpe::BPE;
307 + use tokenizers::pre_tokenizers::byte_level::ByteLevel;
308 use tokenizers::{AddedToken, Tokenizer as HfTokenizer};
309
310 use super::{HuggingFaceTokenizer, Tokenizer};
311
312 + const REGULAR_TOKEN: &str = "<|regular|>";
313 + const SPECIAL_TOKEN: &str = "<|special|>";
314 +
315 fn tiny_bpe_tokenizer() -> HfTokenizer {
316 let vocab = [
317 ("<unk>".to_string(), 0),
@@ -232,6 +338,186 @@ mod tests {
338 HfTokenizer::new(model)
339 }
340
341 + fn ordinary_test_tokenizer_json(fused: bool, with_added_tokens: bool) -> Value {
342 + let mut alphabet: Vec<char> = ByteLevel::alphabet().into_iter().collect();
343 + alphabet.sort_unstable();
344 + let vocab = alphabet
345 + .into_iter()
346 + .enumerate()
347 + .map(|(id, token)| (token.to_string(), json!(id)))
348 + .collect::<serde_json::Map<_, _>>();
349 +
350 + let pre_tokenizer = if fused {
351 + json!({
352 + "type": "Sequence",
353 + "pretokenizers": [
354 + {
355 + "type": "Split",
356 + "pattern": {"Regex": "\\S+|\\s+"},
357 + "behavior": "Isolated",
358 + "invert": false
359 + },
360 + {
361 + "type": "ByteLevel",
362 + "add_prefix_space": false,
363 + "trim_offsets": true,
364 + "use_regex": false
365 + }
366 + ]
367 + })
368 + } else {
369 + json!({
370 + "type": "ByteLevel",
371 + "add_prefix_space": false,
372 + "trim_offsets": true,
373 + "use_regex": true
374 + })
375 + };
376 + let added_tokens = with_added_tokens.then(|| {
377 + json!([
378 + {
379 + "id": 256,
380 + "content": REGULAR_TOKEN,
381 + "single_word": false,
382 + "lstrip": false,
383 + "rstrip": false,
384 + "normalized": true,
385 + "special": false
386 + },
387 + {
388 + "id": 257,
389 + "content": SPECIAL_TOKEN,
390 + "single_word": false,
391 + "lstrip": false,
392 + "rstrip": false,
393 + "normalized": false,
394 + "special": true
395 + }
396 + ])
397 + });
398 +
399 + json!({
400 + "version": "1.0",
401 + "truncation": {
402 + "direction": "Right",
403 + "max_length": 24,
404 + "strategy": "LongestFirst",
405 + "stride": 0
406 + },
407 + "padding": null,
408 + "added_tokens": added_tokens.unwrap_or_else(|| json!([])),
409 + "normalizer": {"type": "NFC"},
410 + "pre_tokenizer": pre_tokenizer,
411 + "post_processor": {
412 + "type": "ByteLevel",
413 + "add_prefix_space": false,
414 + "trim_offsets": true,
415 + "use_regex": true
416 + },
417 + "decoder": {
418 + "type": "ByteLevel",
419 + "add_prefix_space": false,
420 + "trim_offsets": true,
421 + "use_regex": true
422 + },
423 + "model": {
424 + "type": "BPE",
425 + "dropout": null,
426 + "unk_token": null,
427 + "continuing_subword_prefix": null,
428 + "end_of_word_suffix": null,
429 + "fuse_unk": false,
430 + "byte_fallback": false,
431 + "ignore_merges": false,
432 + "vocab": vocab,
433 + "merges": []
434 + }
435 + })
436 + }
437 +
438 + fn write_tokenizer_json(dir: &Path, name: &str, value: &Value) -> PathBuf {
439 + let path = dir.join(name);
440 + std::fs::write(
441 + &path,
442 + serde_json::to_vec(value).expect("serialize tokenizer"),
443 + )
444 + .expect("write tokenizer");
445 + path
446 + }
447 +
448 + fn assert_ordinary_matches_added_empty(
449 + constructor: fn(&Path) -> crate::Result<HuggingFaceTokenizer>,
450 + fused: bool,
451 + ) {
452 + let dir = tempdir().expect("create temp dir");
453 + let added_path = write_tokenizer_json(
454 + dir.path(),
455 + "with-added.json",
456 + &ordinary_test_tokenizer_json(fused, true),
457 + );
458 + let empty_path = write_tokenizer_json(
459 + dir.path(),
460 + "added-empty.json",
461 + &ordinary_test_tokenizer_json(fused, false),
462 + );
463 + let tokenizer = constructor(&added_path).expect("load tokenizer with added tokens");
464 + let added_empty = constructor(&empty_path).expect("load tokenizer with empty added tokens");
465 +
466 + if let super::Backend::Fastokens(inner) | super::Backend::FastokensByteLevel(inner) =
467 + &tokenizer.backend
468 + {
469 + assert_eq!(super::fastokens_fused_split(inner).is_some(), fused);
470 + }
471 +
472 + assert_eq!(
473 + tokenizer.encode(REGULAR_TOKEN, false).unwrap(),
474 + vec![tokenizer.token_to_id(REGULAR_TOKEN).unwrap()]
475 + );
476 + assert_eq!(
477 + tokenizer.encode(SPECIAL_TOKEN, false).unwrap(),
478 + vec![tokenizer.token_to_id(SPECIAL_TOKEN).unwrap()]
479 + );
480 +
481 + for text in [
482 + "",
483 + "hello",
484 + "Cafe\u{301}",
485 + REGULAR_TOKEN,
486 + SPECIAL_TOKEN,
487 + "hello <|regular|> Cafe\u{301} <|special|> tail",
488 + ] {
489 + assert_eq!(
490 + tokenizer.encode_ordinary(text).unwrap(),
491 + added_empty.encode(text, false).unwrap(),
492 + "fused={fused}, text={text:?}",
493 + );
494 + }
495 + if matches!(&tokenizer.backend, super::Backend::Hf(_)) {
496 + assert_eq!(
497 + tokenizer
498 + .encode_ordinary("hello <|regular|> Cafe\u{301} <|special|> tail")
499 + .unwrap()
500 + .len(),
501 + 24,
502 + "HF post-processing must retain configured truncation",
503 + );
504 + }
505 + }
506 +
507 + #[test]
508 + fn hf_ordinary_matches_original_encode_with_added_empty() {
509 + for fused in [false, true] {
510 + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_hf, fused);
511 + }
512 + }
513 +
514 + #[test]
515 + fn fastokens_ordinary_matches_original_encode_with_added_empty() {
516 + for fused in [false, true] {
517 + assert_ordinary_matches_added_empty(HuggingFaceTokenizer::new_fastokens, fused);
518 + }
519 + }
520 +
521 #[test]
522 fn hf_constructor_resolves_added_token_ids() {
523 let mut tokenizer = tiny_bpe_tokenizer();
rust/src/tokenizer/src/incremental.rs
+12
@@ -199,6 +199,10 @@ mod tests {
199 unreachable!()
200 }
201
202 + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> {
203 + unreachable!()
204 + }
205 +
206 fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
207 let bytes = token_ids.iter().map(|id| *id as u8).collect::<Vec<_>>();
208 Ok(String::from_utf8_lossy(&bytes).into_owned())
@@ -273,6 +277,10 @@ mod tests {
277 unreachable!()
278 }
279
280 + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> {
281 + unreachable!()
282 + }
283 +
284 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
285 let mut text = String::new();
286 for &token_id in token_ids {
@@ -410,6 +418,10 @@ mod tests {
418 unreachable!()
419 }
420
421 + fn encode_ordinary(&self, _text: &str) -> Result<Vec<u32>> {
422 + unreachable!()
423 + }
424 +
425 fn decode(&self, token_ids: &[u32], _skip_special_tokens: bool) -> Result<String> {
426 match token_ids {
427 [1] => Ok("abc".into()),
rust/src/tokenizer/src/lib.rs
+4
@@ -25,6 +25,10 @@ pub trait Tokenizer: Send + Sync {
25 /// Encode one prompt string into token IDs.
26 fn encode(&self, text: &str, add_special_tokens: bool) -> Result<Vec<u32>>;
27
28 + /// Equivalent to `encode(text, false)`, except that every added,
29 + /// special, and control-token matcher is bypassed.
30 + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>>;
31 +
32 /// Decode one token sequence into text.
33 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String>;
34
rust/src/tokenizer/src/tekken.rs
+54
@@ -35,6 +35,12 @@ impl Tokenizer for TekkenTokenizer {
35 .map_err(|error| tokenizer_error!("encoding failed: {error}"))
36 }
37
38 + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> {
39 + self.inner
40 + .encode(text, false, false)
41 + .map_err(|error| tokenizer_error!("encoding failed: {error}"))
42 + }
43 +
44 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
45 let policy = if skip_special_tokens {
46 tekken::SpecialTokenPolicy::Ignore
@@ -67,3 +73,51 @@ impl Tokenizer for TekkenTokenizer {
73 self.inner.is_special_token(token_id)
74 }
75 }
76 +
77 +#[cfg(test)]
78 +mod tests {
79 + use base64::Engine as _;
80 + use tekken::config::TokenizerVersion;
81 + use tekken::{SpecialTokenInfo, TokenInfo};
82 +
83 + use super::*;
84 +
85 + fn test_tokenizer() -> TekkenTokenizer {
86 + let vocab = (0_u8..=255)
87 + .map(|byte| TokenInfo {
88 + rank: byte as usize,
89 + token_bytes: base64::engine::general_purpose::STANDARD.encode([byte]),
90 + token_str: None,
91 + })
92 + .collect();
93 + let special_tokens = vec![SpecialTokenInfo {
94 + rank: 0,
95 + token_str: "<control>".to_string(),
96 + is_control: true,
97 + }];
98 + let inner = Tekkenizer::new(
99 + vocab,
100 + &special_tokens,
101 + r"(?s).",
102 + 257,
103 + 1,
104 + TokenizerVersion::V3,
105 + None,
106 + )
107 + .expect("build Tekken tokenizer");
108 + TekkenTokenizer { inner }
109 + }
110 +
111 + #[test]
112 + fn ordinary_matches_tekkens_empty_special_encoding() {
113 + let tokenizer = test_tokenizer();
114 + let text = "user <control> text";
115 + let control_id = tokenizer.token_to_id("<control>").unwrap();
116 + let ordinary_ids = tokenizer.encode_ordinary(text).unwrap();
117 +
118 + assert_eq!(control_id, 0);
119 + assert_eq!(ordinary_ids, tokenizer.encode(text, false).unwrap());
120 + assert!(!ordinary_ids.contains(&control_id));
121 + assert_eq!(tokenizer.decode(&ordinary_ids, false).unwrap(), text);
122 + }
123 +}
rust/src/tokenizer/src/test_utils.rs
+30
@@ -208,6 +208,10 @@ impl Tokenizer for TestTokenizer {
208 Ok(ids)
209 }
210
211 + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> {
212 + Ok(text.as_bytes().iter().copied().map(u32::from).collect())
213 + }
214 +
215 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
216 let mut output = String::new();
217 let mut pending_bytes = Vec::new();
@@ -374,6 +378,32 @@ mod tests {
378 assert!(!tokenizer.is_special_id(0xF002));
379 }
380
381 + #[test]
382 + fn ordinary_encoding_bypasses_all_configured_tokens() {
383 + let tokenizer = TestTokenizer::new()
384 + .with_bos_token("<bos>", 256)
385 + .with_special_token("<control>", 257)
386 + .with_regular_token("<visible>", 258);
387 + let ordinary_text = "user <control> and <visible>";
388 +
389 + assert_eq!(tokenizer.encode("<control>", false).unwrap(), vec![257]);
390 + assert_eq!(tokenizer.encode("<visible>", false).unwrap(), vec![258]);
391 + assert_eq!(
392 + tokenizer.encode_ordinary(ordinary_text).unwrap(),
393 + ordinary_text.as_bytes().iter().copied().map(u32::from).collect::<Vec<_>>()
394 + );
395 +
396 + let mut segmented = tokenizer.encode("<control>", false).unwrap();
397 + segmented.extend(tokenizer.encode_ordinary(ordinary_text).unwrap());
398 + segmented.extend(tokenizer.encode("<visible>", false).unwrap());
399 + assert_eq!(segmented.first(), Some(&257));
400 + assert_eq!(segmented.last(), Some(&258));
401 + assert_eq!(
402 + tokenizer.decode(&segmented, false).unwrap(),
403 + format!("<control>{ordinary_text}<visible>")
404 + );
405 + }
406 +
407 #[test]
408 #[should_panic(expected = "configured test token id 255 overlaps byte fallback range 0..=255")]
409 fn configured_token_id_must_stay_outside_byte_range() {
rust/src/tokenizer/src/tiktoken.rs
+40
@@ -462,6 +462,13 @@ impl Tokenizer for TiktokenTokenizer {
462 })
463 }
464
465 + fn encode_ordinary(&self, text: &str) -> Result<Vec<u32>> {
466 + Ok(match &self.backend {
467 + Backend::Riptoken(backend) => backend.inner.encode_ordinary(text),
468 + Backend::TiktokenRs(backend) => backend.inner.encode_ordinary(text),
469 + })
470 + }
471 +
472 fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result<String> {
473 // Filter passes:
474 //
@@ -752,6 +759,39 @@ mod tests {
759 }
760 }
761
762 + #[test]
763 + fn tiktoken_ordinary_bypasses_every_registered_added_token() {
764 + let dir = tempfile::tempdir().expect("create temp dir");
765 + let bpe_path = write_synthetic_bpe_file(dir.path());
766 + fs::write(
767 + dir.path().join("tokenizer_config.json"),
768 + r#"{
769 + "added_tokens_decoder": {
770 + "257": { "content": "<|im_end|>", "special": true },
771 + "258": { "content": "<|tool_call_begin|>", "special": false }
772 + }
773 + }"#,
774 + )
775 + .expect("write tokenizer_config.json");
776 + fs::write(dir.path().join("config.json"), r#"{"vocab_size": 260}"#)
777 + .expect("write config.json");
778 +
779 + let input = "<|im_end|><|tool_call_begin|><|reserved_token_259|>";
780 + let expected: Vec<u32> = input.as_bytes().iter().copied().map(u32::from).collect();
781 + for backend in explicit_backends(&bpe_path) {
782 + assert_eq!(backend.encode("<|im_end|>", false).unwrap(), vec![257]);
783 + assert_eq!(
784 + backend.encode("<|tool_call_begin|>", false).unwrap(),
785 + vec![258]
786 + );
787 + assert_eq!(
788 + backend.encode("<|reserved_token_259|>", false).unwrap(),
789 + vec![259]
790 + );
791 + assert_eq!(backend.encode_ordinary(input).unwrap(), expected);
792 + }
793 + }
794 +
795 /// `vocab_size` may live under `text_config` for composite (e.g.
796 /// multimodal) configs.
797 #[test]