From bb360a41627deaec77fc6920a42be4d133f8a0ce Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:37:06 +0200 Subject: [PATCH 01/13] perf(cache): serve repeated words from a shared WordCache (#2262) A set-associative cache keyed on the word bytes, plus a scratch pile reused across encode calls, so BPE, WordPiece and Unigram all stop re-tokenizing words they have already seen. Squashed from feat/bpe-cache-min, which also carries #2261 (perf(scratch): reuse Scratch buffers across encode() calls). --- tokenizers/tk-encode/src/models/bpe/model.rs | 113 +- tokenizers/tk-encode/src/models/bpe/word.rs | 2 +- .../tk-encode/src/models/unigram/model.rs | 207 +++- .../tk-encode/src/models/wordpiece/mod.rs | 157 ++- .../tk-encode/src/tokenizer/pipeline.rs | 215 +++- tokenizers/tk-encode/src/utils/cache.rs | 5 +- tokenizers/tk-encode/src/utils/mod.rs | 1 + tokenizers/tk-encode/src/utils/word_cache.rs | 1027 +++++++++++++++++ 8 files changed, 1667 insertions(+), 60 deletions(-) create mode 100644 tokenizers/tk-encode/src/utils/word_cache.rs diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index fa7bcec23f..9e4853e9b7 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -5,6 +5,7 @@ use crate::tokenizer::{Model, Result, Token}; use crate::utils::byte_level::{self}; use crate::utils::cache::{DEFAULT_CACHE_CAPACITY, MAX_LENGTH}; use crate::utils::iter::ResultShunt; +use crate::utils::word_cache::{Lookup, WordCache}; use crate::vocab::bucket_vocab_store::BucketVocabStore; use crate::vocab_store::VocabStore; use ahash::AHashMap; @@ -680,6 +681,8 @@ pub struct PipelineBPE { vocab: BucketVocabStore, merges: MergeMap, ignore_merges: bool, + /// `None` when the tokenizer asked for no cache. + cache_capacity: Option, } enum Atoms { @@ -760,6 +763,7 @@ impl PipelineBPE { ignore_merges, merges, vocab, + cache_capacity: model.cache.map(|c| c.capacity).filter(|&c| c > 0), }) } @@ -827,23 +831,39 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - if self.ignore_merges - && let Some(id) = self.vocab.get_bytes(sequence.as_bytes()) - { - output.push(PipelineToken { id }); - return Ok(()); - } - - // TODO: persistent cache mapping &str -> &[u32] - let BpeScratch { merge_queue, skip, word, + word_cache, } = scratch; - self.merge_word(sequence, merge_queue, skip, word); - output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); + let mut placement = None; + if let Some(cache) = word_cache.as_mut() { + match cache.lookup(sequence.as_bytes()) { + Lookup::Hit(ids) => { + output.extend(ids.iter().map(|&id| PipelineToken { id })); + return Ok(()); + } + Lookup::Miss(at) => placement = at, + } + } + let start = output.len(); + if self.ignore_merges + && let Some(id) = self.vocab.get_bytes(sequence.as_bytes()) + { + output.push(PipelineToken { id }); + } else { + self.merge_word(sequence, merge_queue, skip, word); + output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); + } + // The ids come back out of `output` because that is the only place both + // branches above leave them: `ignore_merges` never touches `word`. + if let Some(cache) = word_cache.as_mut() + && let Some(at) = placement + { + cache.insert(at, output[start..].iter().map(|token| token.id)); + } Ok(()) } @@ -853,6 +873,7 @@ impl pipeline::Model for PipelineBPE { merge_queue: QuaternaryHeap::with_capacity(64), word: Word::with_capacity(64), skip: Vec::new(), + word_cache: self.cache_capacity.map(WordCache::new), } } } @@ -861,6 +882,8 @@ pub struct BpeScratch { pub(crate) merge_queue: QuaternaryHeap, pub(crate) skip: Vec, pub(crate) word: Word, + /// Outlives the encode call that fills it, or it would never see a word twice. + pub(crate) word_cache: Option, } impl ModelScratch for BpeScratch {} @@ -1438,13 +1461,61 @@ mod tests { assert!(pipeline_ids(&pipeline, "").is_empty()); } + // The pool hands the SAME scratch to successive encodes. State left behind by one + // call (an undrained merge queue, a stale word buffer) would corrupt every call + // after it. Drive several inputs through one scratch and check each still matches + // the reference model. + #[test] + fn reused_scratch_matches_fresh() { + let bpe = hello_builder().build().unwrap(); + let reference = bpe.clone(); + let model = PipelineBPE::from_bpe(bpe, false).unwrap(); + let mut scratch = model.init_scratch(); + for input in ["hello", "hell", "helo", "oleh", "hello", "", "hxe"] { + let mut out = Vec::new(); + pipeline::Model::tokenize_pipeline(&model, input, &mut scratch, &mut out).unwrap(); + let got: Vec = out.iter().map(|t| t.id).collect(); + assert_eq!(got, reference_ids(&reference, input), "{input:?}"); + } + } + + // A cache may forget a word, but it must never change one. Run every word twice + // through one scratch, the second time answered from the cache, against a model + // built with no cache at all. + #[test] + fn cached_ids_match_uncached() { + let cached = PipelineBPE::from_bpe(hello_builder().build().unwrap(), false).unwrap(); + let uncached = + PipelineBPE::from_bpe(hello_builder().cache_capacity(0).build().unwrap(), false) + .unwrap(); + + let mut scratch = cached.init_scratch(); + assert!(scratch.word_cache.is_some(), "nothing is being cached"); + for _ in 0..2 { + for word in [ + "hello", + "hell", + "o", + "hellohello", + "hello-a-word-past-fifteen-bytes", + "hxe", + ] { + let mut out = Vec::new(); + pipeline::Model::tokenize_pipeline(&cached, word, &mut scratch, &mut out) + .unwrap(); + let got: Vec = out.iter().map(|t| t.id).collect(); + assert_eq!(got, pipeline_ids(&uncached, word), "{word:?}"); + } + } + } + #[test] fn unknown_char_without_unk_is_dropped() { let bpe = hello_builder().build().unwrap(); let reference = bpe.clone(); let pipeline = PipelineBPE::from_bpe(bpe, false).unwrap(); // 'x' vanishes, making 'h' and 'e' adjacent, so the (h,e) merge - // applies — mirrors the reference model. + // applies, mirroring the reference model. assert_eq!(pipeline_ids(&pipeline, "hxe"), vec![4]); assert_eq!( pipeline_ids(&pipeline, "hxe"), @@ -1615,7 +1686,7 @@ mod tests { /// A gpt2-shaped miniature: the 256 projected single-byte tokens /// (id == byte value) plus `extra` tokens and merges, given in raw - /// space and projected here — like a real byte-level tokenizer.json, + /// space and projected here, like a real byte-level tokenizer.json, /// whose vocab is stored in the projected alphabet. fn byte_level_bpe( extra: &[(&str, u32)], @@ -1672,6 +1743,22 @@ mod tests { ); } + /// A whole-word vocab hit skips the merge loop but not the vocab lookup, + /// and in a byte-level vocab the words that take that path are the + /// commonest ones in the text. Storing the id it found turns the next + /// occurrence into a cache hit. + #[test] + fn ignore_merges_stores_the_whole_word_id() { + let bpe = byte_level_bpe(&[(" hello", 300)], &[], true); + let model = PipelineBPE::from_bpe(bpe, true).unwrap(); + let mut scratch = model.init_scratch(); + let mut out = Vec::new(); + PipelineModel::tokenize_pipeline(&model, " hello", &mut scratch, &mut out).unwrap(); + + let cache = scratch.word_cache.as_mut().unwrap(); + assert_eq!(cache.lookup(b" hello").hit(), Some(&[300u32][..])); + } + #[test] fn byte_level_requires_full_byte_coverage() { // An ASCII-only vocab covers no control/high bytes: building the diff --git a/tokenizers/tk-encode/src/models/bpe/word.rs b/tokenizers/tk-encode/src/models/bpe/word.rs index ac463d6e1a..006aa6a8d5 100644 --- a/tokenizers/tk-encode/src/models/bpe/word.rs +++ b/tokenizers/tk-encode/src/models/bpe/word.rs @@ -276,7 +276,7 @@ impl Word { self.get_chars_iter().collect() } - pub fn get_chars_iter(&self) -> impl Iterator + '_ { + pub fn get_chars_iter(&self) -> impl ExactSizeIterator + '_ { self.symbols.iter().map(|s| s.c) } diff --git a/tokenizers/tk-encode/src/models/unigram/model.rs b/tokenizers/tk-encode/src/models/unigram/model.rs index 1cc3f4b987..7cdd2d18c0 100644 --- a/tokenizers/tk-encode/src/models/unigram/model.rs +++ b/tokenizers/tk-encode/src/models/unigram/model.rs @@ -3,6 +3,7 @@ use super::{ trie::{Trie, TrieBuilder}, }; use crate::utils::cache::{Cache, MAX_LENGTH}; +use crate::utils::word_cache::{Lookup, WordCache}; use crate::vocab_store::VocabStore; use crate::{ pipeline::{self, PipelineToken}, @@ -239,23 +240,33 @@ impl Unigram { if sentence.is_empty() { return Ok(vec![]); } - if self.alpha.is_none() || self.alpha == Some(0.0) { - if let Some(result) = self.cache.get(sentence) { - Ok(result.to_vec()) - } else { - let result = if self.is_optimized { - self.encode_optimized(sentence)? - } else { - self.encode_unoptimized(sentence)? - }; - if sentence.len() < MAX_LENGTH { - self.cache.set(sentence.to_owned(), result.clone()); - } - Ok(result) - } + if self.samples() { + return self.encode_uncached(sentence); + } + if let Some(result) = self.cache.get(sentence) { + return Ok(result.to_vec()); + } + let result = self.encode_uncached(sentence)?; + if sentence.len() < MAX_LENGTH { + self.cache.set(sentence.to_owned(), result.clone()); + } + Ok(result) + } + + /// Whether [`Unigram::alpha`] asks for a tokenization drawn at random from the + /// lattice instead of its best path. Such a result is one draw out of many, so no + /// cache may hold it. + fn samples(&self) -> bool { + matches!(self.alpha, Some(alpha) if alpha != 0.0) + } + + /// What `sentence` encodes to, worked out rather than looked up: the lattice's + /// best path, or a draw from it when [`Unigram::samples`]. + fn encode_uncached(&self, sentence: &str) -> Result> { + if self.is_optimized && !self.samples() { + self.encode_optimized(sentence) } else { - let result = self.encode_unoptimized(sentence)?; - Ok(result) + self.encode_unoptimized(sentence) } } @@ -503,7 +514,10 @@ impl Model for Unigram { } } -pub struct UnigramScratch {} +pub struct UnigramScratch { + /// Outlives the encode call that fills it, or it would never see a word twice. + pub(crate) word_cache: Option, +} impl pipeline::ModelScratch for UnigramScratch {} @@ -511,18 +525,42 @@ impl pipeline::Model for Unigram { type Scratch = UnigramScratch; fn init_scratch(&self) -> Self::Scratch { - Self::Scratch {} + Self::Scratch { + word_cache: match self.cache.capacity { + 0 => None, + capacity => Some(WordCache::new(capacity)), + }, + } } + /// The pipeline asks for ids and nothing else, so this path caches ids, where + /// [`Unigram::encode`] has to hand back the pieces themselves and caches those. A hit + /// skips the lattice, the `String` every piece is built into, and the vocabulary + /// lookup that turns each one back into an id. fn tokenize_pipeline( &self, sequence: &str, - _scratch: &mut Self::Scratch, + scratch: &mut Self::Scratch, output: &mut Vec, ) -> Result<()> { - let str_tokens = self.encode(sequence)?; + if sequence.is_empty() { + return Ok(()); + } + let mut placement = None; + if !self.samples() + && let Some(cache) = scratch.word_cache.as_mut() + { + match cache.lookup(sequence.as_bytes()) { + Lookup::Hit(ids) => { + output.extend(ids.iter().map(|&id| PipelineToken { id })); + return Ok(()); + } + Lookup::Miss(at) => placement = at, + } + } - for string in str_tokens { + let start = output.len(); + for string in self.encode_uncached(sequence)? { match self.token_to_ids.token_to_id(&string) { Some(id) => { output.push(PipelineToken { id }); @@ -547,6 +585,14 @@ impl pipeline::Model for Unigram { } }; } + + // Sampling skipped the lookup, so `placement` is `None` and this stores + // nothing, which is what `samples` says has to happen. + if let Some(cache) = scratch.word_cache.as_mut() + && let Some(at) = placement + { + cache.insert(at, output[start..].iter().map(|token| token.id)); + } Ok(()) } } @@ -710,4 +756,123 @@ mod tests { let tokens = unigram.tokenize("?é").unwrap(); assert_eq!(tokens[0].id, 0); } + + /// Ids 0..=8 are ``, `a`, `b`, `c`, `d`, `cd`, `ab`, `abc`, `abcd`. + fn abcd_vocab() -> Vocab { + vec![ + ("".to_string(), 0.0), + ("a".to_string(), 0.0), + ("b".to_string(), 0.0), + ("c".to_string(), 0.0), + ("d".to_string(), 0.0), + ("cd".to_string(), 1.0), + ("ab".to_string(), 2.0), + ("abc".to_string(), 5.0), + ("abcd".to_string(), 10.0), + ] + } + + fn pipeline_ids(model: &Unigram, sequence: &str, scratch: &mut UnigramScratch) -> Vec { + let mut output = vec![]; + pipeline::Model::tokenize_pipeline(model, sequence, scratch, &mut output).unwrap(); + output.iter().map(|token| token.id).collect() + } + + #[test] + fn pipeline_remembers_what_a_sequence_encoded_to() { + let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap(); + let mut scratch = pipeline::Model::init_scratch(&model); + + let ids = pipeline_ids(&model, "abcd", &mut scratch); + + let cache = scratch + .word_cache + .as_mut() + .expect("Unigram encodes with a cache"); + assert_eq!(cache.lookup(b"abcd").hit(), Some(&ids[..])); + } + + #[test] + fn cache_hits_agree_with_a_cold_run() { + let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap(); + let long = "abcd".repeat(400); + let corpus = [ + "abcdacdxx", + "ab", + // The same sequence again, so this one is served from the cache. + "abcdacdxx", + // Out of the vocabulary, and multibyte. + "東京", + // 1600 bytes, past the longest word the cache will store. + long.as_str(), + "abcdacdxx", + ]; + + let mut warm_scratch = pipeline::Model::init_scratch(&model); + let warm = corpus.map(|sequence| pipeline_ids(&model, sequence, &mut warm_scratch)); + let cold = corpus.map(|sequence| { + let mut scratch = pipeline::Model::init_scratch(&model); + pipeline_ids(&model, sequence, &mut scratch) + }); + + assert_eq!(warm, cold); + } + + #[test] + fn caches_only_the_ids_this_sequence_produced() { + // Every sequence the pipeline hands the model appends to one output buffer, + // so a sequence has to remember its own ids, not everything the buffer holds. + let model = Unigram::from(abcd_vocab(), Some(0), false).unwrap(); + let mut scratch = pipeline::Model::init_scratch(&model); + let mut output = vec![]; + pipeline::Model::tokenize_pipeline(&model, "ab", &mut scratch, &mut output).unwrap(); + pipeline::Model::tokenize_pipeline(&model, "cd", &mut scratch, &mut output).unwrap(); + + let ids: Vec = output.iter().map(|token| token.id).collect(); + assert_eq!(ids, [6, 5]); + let cache = scratch.word_cache.as_mut().unwrap(); + assert_eq!(cache.lookup(b"cd").hit(), Some(&[5u32][..])); + } + + #[test] + fn byte_fallback_ids_survive_the_cache() { + // A piece the vocabulary has no id for becomes one id per byte. The cache + // stores what came out, so a hit has to replay all of them. + let vocab = vec![ + ("".to_string(), 0.0), + ("<0xC3>".to_string(), -0.01), + ("<0xA9>".to_string(), -0.03), + ]; + let model = Unigram::from(vocab, Some(0), true).unwrap(); + let mut scratch = pipeline::Model::init_scratch(&model); + + let ids = pipeline_ids(&model, "é", &mut scratch); + + assert_eq!(ids, [1, 2]); + assert_eq!(pipeline_ids(&model, "é", &mut scratch), ids); + } + + #[test] + fn sampling_is_never_cached() { + // A sampled tokenization is one draw out of many. Remembering it would turn + // every later call on the same text into that same draw. + let mut model = Unigram::from(abcd_vocab(), Some(0), false).unwrap(); + model.alpha = Some(0.5); + let mut scratch = pipeline::Model::init_scratch(&model); + + pipeline_ids(&model, "abcd", &mut scratch); + + let cache = scratch.word_cache.as_mut().unwrap(); + assert_eq!(cache.lookup(b"abcd").hit(), None); + } + + #[test] + fn a_capacity_of_zero_turns_the_cache_off() { + let mut model = Unigram::from(abcd_vocab(), Some(0), false).unwrap(); + model.resize_cache(0); + let mut scratch = pipeline::Model::init_scratch(&model); + + assert_eq!(pipeline_ids(&model, "abcd", &mut scratch), [8]); + assert!(scratch.word_cache.is_none()); + } } diff --git a/tokenizers/tk-encode/src/models/wordpiece/mod.rs b/tokenizers/tk-encode/src/models/wordpiece/mod.rs index a1286fe953..4b9a02dd0a 100644 --- a/tokenizers/tk-encode/src/models/wordpiece/mod.rs +++ b/tokenizers/tk-encode/src/models/wordpiece/mod.rs @@ -4,6 +4,8 @@ use crate::models::bpe::BPE; use crate::pipeline::{self, PipelineToken}; use crate::tokenizer::{Model, Result, Token}; +use crate::utils::cache::DEFAULT_CACHE_CAPACITY; +use crate::utils::word_cache::{Lookup, WordCache}; use ahash::AHashMap; use std::collections::HashMap; use std::convert::TryFrom; @@ -316,6 +318,8 @@ impl Model for WordPiece { pub struct WordPieceScratch { candidate_str: String, + /// Outlives the encode call that fills it, or it would never see a word twice. + word_cache: WordCache, } impl pipeline::ModelScratch for WordPieceScratch {} @@ -353,23 +357,18 @@ impl TryFrom for PipelineWordPiece { } } -impl pipeline::Model for PipelineWordPiece { - type Scratch = WordPieceScratch; - - fn init_scratch(&self) -> Self::Scratch { - Self::Scratch { - candidate_str: String::with_capacity(self.max_input_chars_per_word), - } - } - - fn tokenize_pipeline( +impl PipelineWordPiece { + /// One word, greedily: the longest vocabulary entry it starts with, then the + /// longest entry the rest of it starts with once the continuing-subword prefix + /// is put in front, and so on. A piece with no entry at all anywhere in the + /// word makes the whole word one unk token. + fn tokenize_word( &self, sequence: &str, - scratch: &mut Self::Scratch, - output: &mut Vec, + candidate: &mut String, + output: &mut Vec, ) -> Result<()> { let checkpoint = output.len(); - let candidate = &mut scratch.candidate_str; let char_len = sequence.chars().count(); if char_len > self.max_input_chars_per_word { @@ -411,6 +410,49 @@ impl pipeline::Model for PipelineWordPiece { } } +impl pipeline::Model for PipelineWordPiece { + type Scratch = WordPieceScratch; + + fn init_scratch(&self) -> Self::Scratch { + Self::Scratch { + candidate_str: String::with_capacity(self.max_input_chars_per_word), + word_cache: WordCache::new(DEFAULT_CACHE_CAPACITY), + } + } + + /// A hit skips `tokenize_word`: one trie search per piece of the word, + /// each over a fresh copy of what is left to match. + fn tokenize_pipeline( + &self, + sequence: &str, + scratch: &mut Self::Scratch, + output: &mut Vec, + ) -> Result<()> { + if sequence.is_empty() { + return Ok(()); + } + let WordPieceScratch { + candidate_str, + word_cache, + } = scratch; + + let placement = match word_cache.lookup(sequence.as_bytes()) { + Lookup::Hit(ids) => { + output.extend(ids.iter().map(|&id| PipelineToken { id })); + return Ok(()); + } + Lookup::Miss(at) => at, + }; + + let start = output.len(); + self.tokenize_word(sequence, candidate_str, output)?; + if let Some(at) = placement { + word_cache.insert(at, output[start..].iter().map(|token| token.id)); + } + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -419,4 +461,93 @@ mod tests { fn test_error_display() { assert!(format!("{}", Error::MissingUnkToken).contains("Missing [UNK] token")); } + + /// `hello` is in the vocabulary whole and as `hell` + `##o`, so the + /// longest-match walk has something to choose; `world` gives a second + /// one-token word. + fn pipeline_wordpiece() -> PipelineWordPiece { + let vocab: Vocab = [ + ("[UNK]", 0u32), + ("hell", 1), + ("##o", 2), + ("hello", 3), + ("world", 4), + ] + .into_iter() + .map(|(token, id)| (token.to_string(), id)) + .collect(); + let model = WordPiece::builder() + .vocab(vocab) + .max_input_chars_per_word(8) + .build() + .unwrap(); + PipelineWordPiece::try_from(model).unwrap() + } + + fn pipeline_ids( + model: &PipelineWordPiece, + sequence: &str, + scratch: &mut WordPieceScratch, + ) -> Vec { + let mut output = vec![]; + pipeline::Model::tokenize_pipeline(model, sequence, scratch, &mut output).unwrap(); + output.iter().map(|token| token.id).collect() + } + + #[test] + fn pipeline_remembers_what_a_word_encoded_to() { + let model = pipeline_wordpiece(); + let mut scratch = pipeline::Model::init_scratch(&model); + + let ids = pipeline_ids(&model, "hello", &mut scratch); + + assert_eq!(scratch.word_cache.lookup(b"hello").hit(), Some(&ids[..])); + } + + #[test] + fn cache_hits_agree_with_a_cold_run() { + let model = pipeline_wordpiece(); + let long = "hello".repeat(300); + let corpus = [ + "hello", + // No id for `##w`, so the whole word is one unk token. + "hellow", + // Both again, so these two are served from the cache. + "hello", + "hellow", + // Past `max_input_chars_per_word`, which is another unk token. + "hellohello", + "hellohello", + // Out of the vocabulary, and multibyte. + "東京", + "東京", + // 1500 bytes, past the longest word the cache will store. + long.as_str(), + long.as_str(), + ]; + + let mut warm_scratch = pipeline::Model::init_scratch(&model); + let warm = corpus.map(|sequence| pipeline_ids(&model, sequence, &mut warm_scratch)); + let cold = corpus.map(|sequence| { + let mut scratch = pipeline::Model::init_scratch(&model); + pipeline_ids(&model, sequence, &mut scratch) + }); + + assert_eq!(warm, cold); + } + + #[test] + fn caches_only_the_ids_this_word_produced() { + // Every word the pipeline hands the model appends to one output buffer, + // so a word has to remember its own ids, not everything the buffer holds. + let model = pipeline_wordpiece(); + let mut scratch = pipeline::Model::init_scratch(&model); + let mut output = vec![]; + pipeline::Model::tokenize_pipeline(&model, "hello", &mut scratch, &mut output).unwrap(); + pipeline::Model::tokenize_pipeline(&model, "world", &mut scratch, &mut output).unwrap(); + + let ids: Vec = output.iter().map(|token| token.id).collect(); + assert_eq!(ids, [3, 4]); + assert_eq!(scratch.word_cache.lookup(b"world").hit(), Some(&[4u32][..])); + } } diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 16b7430b78..01e2172961 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1,5 +1,7 @@ use std::cell::RefCell; use std::convert::TryInto; +use std::mem; +use std::sync::{Mutex, PoisonError}; use std::{borrow::Cow, convert::TryFrom}; use atomsplit::classify::classify; @@ -56,7 +58,7 @@ pub(crate) fn classify_into_spans( } classify(bytes, &mut tags[..n]); let k = fsm(bytes, &tags[..n], &mut spans[..n + 1]); - out.extend_from_slice(&spans[..k]); // same type now — plain memcpy, no per-token conversion + out.extend_from_slice(&spans[..k]); // same type now: plain memcpy, no per-token conversion }); } @@ -323,7 +325,7 @@ impl TryFrom<&PostProcessorWrapper> for PipelinePostProcessor { } } -/// An output token. Carries only the vocabulary `id` — offsets and the token +/// An output token. Carries only the vocabulary `id`, since offsets and the token /// string are dropped, which is all an encode-only caller needs. #[derive(Debug, Clone, Copy)] pub struct PipelineToken { @@ -448,6 +450,78 @@ pub struct PipelineTokenizer { pre_tokenizer: PipelinePreTokenizer, model: PipelineModel, post_processor: PipelinePostProcessor, + scratch_pool: ScratchPool, +} + +/// A pool of [`PipelineModelScratch`]. +/// +/// When calling [`PipelineTokenizer::encode`], an instance of [`PipelineModelScratch`] is taken out of this pool +/// and given to the tokenizer. When the encoding is done, the scratch buffer is returned to the pool and can be +/// reused by later calls. +/// +/// The reusability matters because the scratch buffer may hold cache structures which are more useful when reused, +/// and less importantly it saves an extra allocation for an fresh buffer every time. +struct ScratchPool(Mutex>); + +impl ScratchPool { + fn new() -> Self { + Self(Mutex::new(Vec::new())) + } + + /// Get a scratch buffer from the pool, wrapped in a [`ScratchGuard`]. + /// When the [`ScratchGuard`] gets dropped, the scratch buffer is pushed back to the pool. + fn get<'a>(&'a self, model: &PipelineModel) -> ScratchGuard<'a> { + // The Mutex lock is held just long enough to pop the scratch out of the pool + let taken = self.0.lock().unwrap_or_else(PoisonError::into_inner).pop(); + ScratchGuard { + // If there was no scratch buffer available in the pool, we build.a fresh one + scratch: taken.unwrap_or_else(|| model.init_scratch()), + pool: self, + } + } + + #[cfg(test)] + fn len(&self) -> usize { + self.0.lock().unwrap_or_else(PoisonError::into_inner).len() + } +} + +/// A wrapper around [`PipelineModelScratch`]. +/// Implements [`Deref`] and [`DerefMut`], so it behaves as [`PipelineModelScratch`]. +/// +/// When it gets dropped, it pushes [`Self::scratch`] back into the shared [`Self::pool`] so it can +/// get reused by a later call to [`PipelineTokenizer::encode`]. +/// +/// TODO @McPatate : The Mutex can create contention, to be replaced by a better access pattern +struct ScratchGuard<'a> { + scratch: PipelineModelScratch, + pool: &'a ScratchPool, +} + +impl Drop for ScratchGuard<'_> { + fn drop(&mut self) { + // Steals the scratch buffer from self, replaces it with PipelineModelScratch::default() + let scratch = mem::take(&mut self.scratch); + // Push the scratch back in the pool + self.pool + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(scratch); + } +} + +impl std::ops::Deref for ScratchGuard<'_> { + type Target = PipelineModelScratch; + fn deref(&self) -> &PipelineModelScratch { + &self.scratch + } +} + +impl std::ops::DerefMut for ScratchGuard<'_> { + fn deref_mut(&mut self) -> &mut PipelineModelScratch { + &mut self.scratch + } } impl TryFrom<&Tokenizer> for PipelineTokenizer { @@ -569,6 +643,7 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { .map(PipelinePostProcessor::try_from) .transpose()? .unwrap_or_default(), + scratch_pool: ScratchPool::new(), }) } } @@ -599,7 +674,7 @@ impl PipelineTokenizer { pub fn encode(&self, input: &str, add_special_tokens: bool) -> Result> { let mut output = Vec::new(); let mut pre_tokens = Vec::new(); - let mut scratch = self.model.init_scratch(); + let mut scratch = self.scratch_pool.get(&self.model); self.encode_generic::<{ Self::STAGE_POSTPROCESS }>( input, @@ -613,7 +688,7 @@ impl PipelineTokenizer { /// Decode token ids back to a `String`. /// - /// Not implemented yet — the pipeline decode path is being built. It fails + /// Not implemented yet: the pipeline decode path is being built. It fails /// loud (rather than returning a plausible-but-wrong string) so the oracle /// test and the comparative benchmark report decode as *pending* instead of /// silently validating garbage. Implementing this flips the ignored @@ -623,7 +698,7 @@ impl PipelineTokenizer { } /// Decode several id sequences at once, one `String` per input. Mirrors the - /// released `decode_batch`; sequential (KISS) — behavior-identical to a + /// released `decode_batch`; sequential (KISS), behavior-identical to a /// parallel map, since each [`decode`](Self::decode) is independent. pub fn decode_batch( &self, @@ -638,7 +713,7 @@ impl PipelineTokenizer { /// Incremental decode: feed ids one at a time via [`PipelineDecodeStream::step`]. /// Same prefix-tracking scheme as the released `DecodeStream`, built on - /// [`decode`](Self::decode) — so it is correct exactly where `decode` is. + /// [`decode`](Self::decode), so it is correct exactly where `decode` is. pub fn decode_stream(&self, skip_special_tokens: bool) -> PipelineDecodeStream<'_> { PipelineDecodeStream { tokenizer: self, @@ -651,7 +726,7 @@ impl PipelineTokenizer { /// Single source of truth for the encode pipeline, generic over how many stages /// run. `STAGE` is a **const generic**, so `if STAGE >= …` folds at compile time and - /// the disabled stages are compiled out — the full specialization + /// the disabled stages are compiled out, so the full specialization /// ([`STAGE_POSTPROCESS`], which [`encode`](Self::encode) calls) is branchless and /// identical to a hand-written full pipeline, while the benchmark drives lower /// `STAGE` values to time each stage's marginal cost (the ablation ladder), e.g. @@ -660,7 +735,7 @@ impl PipelineTokenizer { /// [`STAGE_POSTPROCESS`]: Self::STAGE_POSTPROCESS /// /// `output` and the `pre_tokens` scratch are caller-owned so a benchmark can reuse - /// them across calls and observe both buffers to anchor the ablation levels — the + /// them across calls and observe both buffers to anchor the ablation levels. The /// library itself stays free of any `black_box`/timing artifact. #[doc(hidden)] // public only so `examples/fixture_bench.rs` can drive partial stages pub fn encode_generic( @@ -785,7 +860,7 @@ pub enum SplitPolicy { Isolate, } -/// Splits `text` into same-class groups, emitting each as a [`Split`] +/// Splits `text` into same-class groups, emitting each as a [`Span`] /// according to its [`SplitPolicy`]. /// /// `classify` maps each char to a small `Copy + Eq` class, the current @@ -909,7 +984,7 @@ pub fn split_matches( ) { use SplitDelimiterBehavior::*; - // (offsets, should_remove) — mirrors `NormalizedString::split`. + // (offsets, should_remove), mirroring `NormalizedString::split`. let splits: Vec<((usize, usize), bool)> = match behavior { Isolated => matches.into_iter().map(|(o, _)| (o, false)).collect(), Removed => matches, // should_remove == is_match @@ -1045,11 +1120,19 @@ impl Model for PipelineModel { } } +/// A set of buffers and other state the model needs to encode efficiently, +/// reused among calls to [`PipelineTokenizer::encode`]. +/// +/// Each model gets its own variant. +#[derive(Default)] pub enum PipelineModelScratch { BPE(BpeScratch), WordLevel(()), WordPiece(WordPieceScratch), Unigram(UnigramScratch), + /// We need a default value to be able to use [`mem::take`] in [`ScratchGuard::drop`] + #[default] + None, } impl ModelScratch for PipelineModelScratch {} @@ -1415,4 +1498,116 @@ mod tests { let err = conversion_error(&tok); assert!(err.contains("not supported"), "{}", err); } + + /// A BPE pipeline that merges "hello" into the single id 7. + fn hello_pipeline() -> PipelineTokenizer { + use crate::models::bpe::{BpeBuilder, Merges, Vocab}; + + let vocab: Vocab = [ + ("h", 0u32), + ("e", 1), + ("l", 2), + ("o", 3), + ("he", 4), + ("hel", 5), + ("hell", 6), + ("hello", 7), + ] + .into_iter() + .map(|(s, i)| (s.to_string(), i)) + .collect(); + let merges: Merges = vec![ + ("h".to_string(), "e".to_string()), + ("he".to_string(), "l".to_string()), + ("hel".to_string(), "l".to_string()), + ("hell".to_string(), "o".to_string()), + ]; + let bpe = BpeBuilder::default() + .vocab_and_merges(vocab, merges) + .build() + .unwrap(); + PipelineTokenizer::try_from(&Tokenizer::new(bpe)).unwrap() + } + + // The pool exists so ONE `&self` tokenizer can be shared across rayon workers. Encode + // the same input from thousands of threads through a single instance; each must get a + // private scratch and produce the sequential result. Two threads sharing a scratch + // would corrupt some of them. This only compiles if `PipelineTokenizer: Sync`, + // which the pool has to preserve. + #[test] + fn encode_shared_across_threads() { + use rayon::prelude::*; + + let pipeline = hello_pipeline(); + + let want: Vec = pipeline + .encode("hello", false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want, vec![7]); + + let all_match = (0..10_000u32).into_par_iter().all(|_| { + pipeline + .encode("hello", false) + .unwrap() + .iter() + .map(|t| t.id) + .collect::>() + == want + }); + assert!(all_match); + } + + // Reusing scratches is the whole point of the pool, so it must not build one per call: + // one thread encoding in a loop has to keep coming back to the same scratch, and a + // burst of N threads must leave at most N behind for later calls to use. + #[test] + fn scratches_are_reused_rather_than_piling_up() { + use std::sync::Barrier; + + let pipeline = hello_pipeline(); + for _ in 0..1000 { + pipeline.encode("hello", false).unwrap(); + } + assert_eq!(pipeline.scratch_pool.len(), 1); + + let threads = 64; + let all_holding = Barrier::new(threads); + std::thread::scope(|scope| { + for _ in 0..threads { + scope.spawn(|| { + let scratch = pipeline.scratch_pool.get(&pipeline.model); + all_holding.wait(); + drop(scratch); + }); + } + }); + + let after_burst = pipeline.scratch_pool.len(); + assert!( + after_burst <= threads, + "{after_burst} scratches kept for {threads} threads" + ); + for _ in 0..1000 { + pipeline.encode("hello", false).unwrap(); + } + assert_eq!(pipeline.scratch_pool.len(), after_burst); + } + + // A scratch coming back out of the pool has to still know the words of the last + // encode: a cache emptied between calls would never hit. + #[test] + fn the_word_cache_outlives_the_encode_call() { + let pipeline = hello_pipeline(); + pipeline.encode("hello", false).unwrap(); + + let mut scratch = pipeline.scratch_pool.get(&pipeline.model); + let PipelineModelScratch::BPE(bpe) = &mut *scratch else { + panic!("a BPE pipeline encodes with a BPE scratch"); + }; + let cache = bpe.word_cache.as_mut().expect("BPE encodes with a cache"); + assert_eq!(cache.lookup(b"hello").hit(), Some(&[7u32][..])); + } } diff --git a/tokenizers/tk-encode/src/utils/cache.rs b/tokenizers/tk-encode/src/utils/cache.rs index 15c6b65f18..4417e9847e 100644 --- a/tokenizers/tk-encode/src/utils/cache.rs +++ b/tokenizers/tk-encode/src/utils/cache.rs @@ -3,8 +3,9 @@ use std::borrow::Borrow; use std::hash::Hash; use std::sync::RwLock; -/// The default capacity for a `BPE`'s internal cache. -pub static DEFAULT_CACHE_CAPACITY: usize = 10_000; +/// The default capacity of a model's cache, whether it is a [`Cache`] or a +/// [`WordCache`](crate::utils::word_cache::WordCache). +pub static DEFAULT_CACHE_CAPACITY: usize = 65_536; /// The maximum length we should cache in a model /// Strings that are too long have minimal chances to cache hit anyway pub static MAX_LENGTH: usize = 256; diff --git a/tokenizers/tk-encode/src/utils/mod.rs b/tokenizers/tk-encode/src/utils/mod.rs index 4003f9e097..56e97dffdc 100644 --- a/tokenizers/tk-encode/src/utils/mod.rs +++ b/tokenizers/tk-encode/src/utils/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod cache; #[cfg(feature = "http")] pub(crate) mod from_pretrained; +pub(crate) mod word_cache; // Optional system-regex backend, needed only for a *regex* pattern that atomsplit does not cover. // With `fancy-regex` off a stub compiles and those patterns error at load. Everything else works diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs new file mode 100644 index 0000000000..96add857dc --- /dev/null +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -0,0 +1,1027 @@ +//! A table that remembers which token ids a word encodes to, so a model only has +//! to work it out once. +//! +//! # Context +//! +//! A tokenizer turns text into *token ids*: small integers a language model reads. +//! It gets there in three stages: +//! - It **normalizes** the text (for example: lowercasing), +//! - It **pre-tokenizes** the normalized text (ie, cuts it into small pieces, usually words or fragments of words), +//! - It runs a tokenization **model** over each piece to produce token ids +//! +//! This module calls those pieces *words*, because that is what they nearly always are. +//! +//! ```text +//! "The cat sat on The Mat" +//! │ +//! │ normalize, then pre-tokenize +//! ▼ +//! "the" "cat" "sat" "on" "the" "mat" ← words +//! │ │ │ │ │ │ +//! ▼ ▼ ▼ ▼ ▼ ▼ +//! ┌─────────────────────────────────────────────┐ +//! │ Tokenization Model │ +//! └─────────────────────────────────────────────┘ +//! │ │ │ │ │ │ +//! ▼ ▼ ▼ ▼ ▼ ▼ +//! [12] [87] [43] [9] [12] [64] ← token ids, made up here +//! ``` +//! +//! The model stage is expensive, due to the underlying algorithms: +//! +//! - [**BPE**](crate::models::bpe::BPE) (byte-pair encoding) starts from the word's individual +//! bytes and repeatedly glues the best-ranked neighboring pair together, until no pair +//! left in the word can be merged. Every merge changes the word and the search +//! for the next best pair starts again ([`Word::merge_all`](crate::models::bpe::Word::merge_all)). +//! - [**Unigram**](crate::models::unigram::Unigram) considers the many ways of cutting the word +//! into pieces the vocabulary knows, scores them, and keeps the best. That is a search over a +//! lattice of candidate cuts ([`Lattice::viterbi`](crate::models::unigram::Lattice::viterbi)). +//! - [**WordPiece**](crate::models::wordpiece::WordPiece) walks the word from the front, taking +//! the longest piece the vocabulary holds, then carries on from where that piece ended. That is +//! one search per piece of the word. +//! +//! Two things make that expense avoidable: +//! - Words repeat, in every kind of text. +//! - A word always encodes to the same ids. +//! +//! So a word only ever has to go through the model once. Cache the ids that came out, and +//! every later occurrence of that word skips the model altogether. +//! +//! This module is the table that does the remembering, in bounded memory: +//! +//! ```text +//! "The cat sat on The Mat" +//! │ +//! │ normalize, then pre-tokenize +//! ▼ +//! "the" "cat" "sat" "on" "the" "mat" ← words, one at a time +//! │ │ │ │ │ │ +//! ▼ ▼ ▼ ▼ ▼ ▼ +//! ┌─────────────────────────────────────────────┐ +//! │ WordCache: have I encoded this word yet? │ +//! └─────────────────────────────────────────────┘ +//! │ │ +//! miss │ │ hit +//! ▼ │ +//! ┌─────────────────────┐ │ +//! │ Model (expensive) │ │ +//! └─────────────────────┘ │ +//! │ │ +//! │ store the ids │ the ids, straight from the table +//! └─────┬──────────┘ +//! ┌──────┬──────┼──────┬──────┬──────┐ +//! ▼ ▼ ▼ ▼ ▼ ▼ +//! [12] [87] [43] [9] [12] [64] ← the same ids, most of them for free +//! ▲ +//! └ the second "the": a cache hit +//! ``` +//! +//! # What the cache is allowed to get wrong +//! +//! The cache is free to *forget*: if a word's ids are gone, the model works them +//! out again. What it must never do is hand back the *wrong* ids. Every +//! trade-off below spends the first freedom and none of the second. +//! +//! # The table +//! +//! The cache is one long row of numbered **slots**. Each slot holds one word and +//! the ids that word encodes to. The number of slots is fixed when the cache is +//! built and never changes, which is what keeps the memory bounded. +//! +//! To decide where a word belongs, the cache turns the word into a number, its +//! **hash**. Two different pieces of that number are used, and they never overlap: +//! +//! ```text +//! "the" +//! │ +//! │ hash +//! ▼ +//! ┌─────────────────────────────────────────┐ +//! │ 1010 0111 ................... 0000 0101 │ +//! └────┬───────────────────────────────┬────┘ +//! │ the top 8 bits │ the bottom bits +//! ▼ ▼ +//! the tag: A7 the home slot: 5 +//! ``` +//! +//! The **home slot** is where the word would like to live. The **tag** is a +//! one-byte summary of the word. Tags are kept in their own row, one byte per +//! slot, beside the row of entries: +//! +//! ```text +//! slot index: 0 1 2 3 4 5 6 7 ... +//! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ +//! tags │ · │ 9B │ · │ · │ C4 │ A7 │ 31 │ A7 │ 1 byte each +//! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ +//! entries │ │"of" │ │ │"cat"│"the"│"sat"│"hat"│ 32 bytes each +//! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ +//! ▲ +//! home slot of "the", +//! and of "hat" as well +//! +//! · = empty slot +//! ``` +//! +//! A word does not always get its home slot: another word may have taken it +//! first, which is what happened to "hat" above. So the word may be in its home +//! slot or in any of the fifteen slots after it. Those sixteen slots are the +//! word's **window**, and the word is either in there or nowhere +//! ([`WALK_WINDOW`]). +//! +//! # Looking a word up +//! +//! [`WordCache::lookup`] walks the window one tag at a time, starting at the home +//! slot, and answers both of a lookup's questions on the way: which slot holds +//! this word, and where would this word go if none of them does. +//! +//! ```text +//! looking up "hat" (tag A7, home slot 5) +//! +//! slot: 5 6 7 8 9 10 11 12 ... +//! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ +//! tags │ A7 │ 31 │ A7 │ · │ 5F │ A7 │ · │ · │ +//! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ +//! entries │"the"│"sat"│"hat"│ │"on" │"mat"│ │ │ +//! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ +//! ▲ ▲ ▲ ▲ +//! │ │ │ └ tagged A7 too, but past the +//! │ │ │ empty slot: out of reach +//! │ │ └ empty: the search stops here +//! │ └ tag matches, and the key says "hat": found it +//! └ tag matches, but the key says "the", so keep looking +//! ``` +//! +//! The walk reads tags, not entries: ruling a slot out costs one byte of memory +//! traffic rather than the 32 an entry takes. It rarely gets far. On running text a +//! word is usually in its home slot, or the slot after it, or not in the table at +//! all. +//! +//! A matching tag is a hint, not an answer. A tag is only one byte of the hash, so +//! about one slot in 255 matches a word that is not there, which is exactly what +//! slot 5 does above. Each matching slot is therefore read and its **key** +//! compared. The key is the copy of the word kept inside the slot, and it is what +//! settles the question. +//! +//! The search stops at the first empty slot. That is safe because storing stops +//! there too: a word is always put in the first empty slot of its window, so no +//! word can ever sit past one. This is why slot 10 is out of reach: whatever is +//! in it, it is not "hat", because "hat" would have taken slot 8. +//! +//! A lookup ends one of two ways ([`Lookup`]): +//! +//! - **Hit**: the ids, ready to use. Nothing is written: a hit leaves the table +//! exactly as it was. +//! - **Miss**: a [`Placement`], the slot this word should take. The caller runs +//! the model, then hands the placement back to [`WordCache::insert`] along with +//! the ids. The placement carries the hash work the lookup already did, so +//! storing the word costs no second hash and no second read of the tags. +//! +//! # Storing a word +//! +//! [`WordCache::insert`] takes the placement from the miss and the ids the model +//! produced, and does four things: +//! +//! 1. **Builds the entry.** The word and its ids go in the slot if they fit; what +//! does not fit is copied into an overflow buffer (below). If those buffers are +//! full, the insert is dropped and nothing changes, since the cache is allowed +//! to forget. +//! 2. **Clears the old entry, if there was one.** When the window had no empty +//! slot the placement points at the home slot, and that word's ids are lost. +//! Any overflow space it was using is handed back. +//! 3. **Writes the tag**, which is what makes the slot findable. +//! 4. **Writes the entry.** +//! +//! # What a slot holds +//! +//! A slot is 32 bytes: sixteen for the key, twelve for the ids, one that says how +//! many ids there are, and three of padding. It comes in three shapes +//! ([`CachedWord`]). The first is the common case, and the reason the sizes are what +//! they are: the whole answer (word and ids) is right there, so a hit reads one +//! slot and stops. +//! +//! ```text +//! ├───── key ─────┤├─── ids ───┤├─────┤ +//! how many +//! +//! "the" → [464] ┌────────────────┬────────────┬─────┐ +//! short word, up to 3 ids │ t h e 3 │ 464 │ 1 │ +//! └────────────────┴────────────┴─────┘ +//! the word, then the ids +//! its length themselves +//! +//! "hello" → [40,71,12,9,33] ┌────────────────┬────────────┬─────┐ +//! short word, too many ids │ h e l l o 5 │ ids are at │ ▒ │ +//! └────────────────┴────────────┴─────┘ +//! +//! "counterrevolutionary" → … ┌────────────────┬────────────┬─────┐ +//! word too long for the key │ hash of word H │ both are at│ ▒ │ +//! └────────────────┴────────────┴─────┘ +//! H says "a hash, where to look in +//! not the word" the buffers +//! +//! ▒ = "not in the slot, follow the offsets" +//! ``` +//! +//! A word of up to 15 bytes is kept whole, with its length in the sixteenth byte. +//! A longer word cannot fit, so the slot keeps the word's hash instead and the +//! bytes go to a buffer. Two different long words can hash to the same number, so +//! for those, and only those, a match is confirmed by comparing the bytes. +//! +//! The two overflow buffers ([`Arena`]) grow up to a budget and then stop. When an +//! entry is evicted, the space it used goes on a free list for its exact +//! length, so the next word of that shape reuses it. Nothing is ever compacted or +//! moved. Words longer than [`MAX_WORD_BYTES`] are not cached at all. +//! +//! # Which entry gets evicted +//! +//! When a word's window is full, one of its entries gets evicted, and it is whatever +//! sits in the home slot. Nothing is measured and nothing is ranked: the newcomer +//! takes the slot it wanted in the first place. +//! +//! A cleverer rule is possible, and this module used to have one: a use counter per +//! entry, the least-used slot of the window loses it, and the counters fade so that +//! a word that was busy early in a document cannot keep its slot forever. Measured +//! against the plain rule, it did keep slightly more of the words worth keeping, and +//! it lost more than it gained to the counter that every hit had to write, so it +//! went. Two things let the plain rule get away with it: nothing is evicted until +//! all sixteen slots of a window are taken, and words repeat closely enough that a +//! word's next use usually comes round before anything has had the chance to evict +//! it. +//! +//! # Why it is built this way +//! +//! - **Tags live in their own row** so that a walk over sixteen slots reads sixteen +//! bytes, which is one or two cache lines. The entries are 32 bytes each, so the +//! same sixteen slots would be 512 bytes and eight cache lines. +//! - **A tag is one byte of the hash**, so a slot that cannot hold the word is ruled +//! out without reading it. All eight bits carry hash: emptiness needs a value no +//! live tag can take rather than a bit of its own, and one value is cheaper than +//! one bit ([`make_tag`]). +//! - **The tag uses the top bits of the hash** because the bottom bits already +//! chose the home slot, and every slot in a window would share those. +//! - **A window is sixteen slots**: how far a word may end up from its home slot +//! before the cache stops looking and evicts instead. Long enough that a full +//! window is uncommon, short enough that walking a full one stays cheap. +//! - **The key is 128 bits** because 15 bytes of word plus one byte of length +//! covers nearly every word whole, and comparing it is then one comparison of one +//! value instead of following a pointer to bytes elsewhere in memory. +//! - **The length is part of the key**, or `"a"` and `"a\0"` would look identical. +//! - **Up to three ids fit in the slot** because most words encode to one, two or +//! three, and that makes a hit a single read. +//! - **The number of slots is a power of two**, so picking the home slot is one +//! bit operation rather than a division. So is folding a step of the walk back +//! into the table when a window runs off the end. +//! - **A miss hands back a placement** because the lookup already knows the hash, +//! the tag and the slot to use, and making the insert work them out again would +//! double that cost on every new word. +//! - **Freed overflow space is filed by exact length**: every run is at most +//! [`MAX_WORD_BYTES`] long, so there can be one free list per length and a freed +//! run always fits the next word of that shape exactly. +//! +//! # Where the ideas come from +//! +//! - [Swiss Tables] is where the control byte comes from: one byte per slot, holding +//! a flag bit and seven bits of hash. Here the whole byte is hash, because the flag +//! bit was there to be tested sixteen lanes at a time and this walk tests one. +//! - [gigatoken] is a BPE tokenizer with a pre-token cache built from the same +//! parts: `u128` packed keys with the length in the top byte, self-contained +//! 32-byte entries, ids inline. It never evicts (it doubles at 3/4 load) and +//! leans on huge pages and prefetching, because its table is sized for DRAM +//! rather than for a CPU cache. +//! - [TinyLFU] (Einziger, Friedman & Manes) is the use-counter-and-fade rule this +//! module tried for eviction and dropped, as its reference implementation +//! [Caffeine] does it. Kept here for whoever wants to try it again. +//! - [huggingface/tokenizers#2234] is an open-addressed cache for this same encode +//! pipeline, arrived at in parallel, fused into the pre-tokenizer's split loop. +//! +//! [Swiss Tables]: https://abseil.io/about/design/swisstables +//! [gigatoken]: https://github.com/marcelroed/gigatoken +//! [TinyLFU]: https://arxiv.org/abs/1512.00727 +//! [Caffeine]: https://github.com/ben-manes/caffeine/blob/master/caffeine/src/main/java/com/github/benmanes/caffeine/cache/FrequencySketch.java +//! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 + +use ahash::RandomState; + +// ---------------------------------------------------------------- the cache + +/// Longest word the cache will store, in bytes. +const MAX_WORD_BYTES: usize = 1024; + +/// Word bytes to token ids. See the module docs for the design. +pub struct WordCache { + /// The slots, one word each. Their number is a power of two, so a word's home + /// slot is just the bottom bits of its hash: + /// ```text + /// hash(word) & index_mask + /// ``` + cached_words: Box<[CachedWord]>, + + /// Hashes a slot's key: the packed word when it is short enough, the word's + /// bytes when it is not. See [`WordCache::make_word_key`]. + hasher: RandomState, + + /// `cached_words.len() - 1`. Masks a hash down to its bottom bits, which give the + /// word's home slot. + index_mask: usize, + + /// One tag per slot, at the slot's own index: the top byte of the word's hash + /// ([`make_tag`]), or [`EMPTY`] when the slot holds nothing. + /// + /// A walk reads this row instead of the entries, so ruling a slot out costs one + /// byte rather than the 32 an entry takes. + tags: Box<[u8]>, + + /// Holds the word's bytes when they don't fit in a [`CachedWord`] key. + word_bytes_arena: Arena, + + /// Holds the word's ids when they don't fit in a [`CachedWord`]. + token_ids_arena: Arena, +} + +impl WordCache { + /// `capacity` is rounded up to a power of two, and to at least one full window. + pub fn new(capacity: usize) -> Self { + let n_slots = capacity.next_power_of_two().max(WALK_WINDOW); + Self { + hasher: RandomState::new(), + cached_words: vec![CachedWord::default(); n_slots].into_boxed_slice(), + tags: vec![EMPTY; n_slots].into_boxed_slice(), + word_bytes_arena: Arena::new(n_slots * 48), + token_ids_arena: Arena::new(n_slots * 16), + index_mask: n_slots - 1, + } + } + + /// The ids `word` encoded to last time, or the [`Placement`] it should be stored + /// in once the model has worked them out. See [`Lookup`]. + pub fn lookup<'c, 'w>(&'c self, word: &'w [u8]) -> Lookup<'c, 'w> { + if word.len() > MAX_WORD_BYTES { + return Lookup::Miss(None); + } + + let (key, hash) = self.make_word_key(word); + match self.find_word_in_cache(key, hash, word) { + Walk::Found(index) => { + let slot = self.cached_words[index]; + Lookup::Hit(if slot.ids_stored_in_arena() { + self.token_ids_arena.get(slot.ids_off(), slot.ids_len()) + } else { + &self.cached_words[index].word_ids[..slot.inline_id_count as usize] + }) + } + Walk::Absent(placement) => Lookup::Miss(Some(placement)), + } + } + + /// Store `ids` as the encoding of the word that `at` was built for. + /// + /// Overwrites whatever the slot held, which only happens when the word's window + /// was full. The module docs say why the choice does not have to be cleverer than + /// that. If the arenas have no room for the entry, nothing is stored. + pub fn insert(&mut self, at: Placement<'_>, ids: impl ExactSizeIterator) { + let Some(cached_word) = self.build_entry(at.key, at.word, ids) else { + return; + }; + if self.tags[at.index] != EMPTY { + self.reclaim(at.index); + } + self.tags[at.index] = at.tag; + self.cached_words[at.index] = cached_word; + } + + /// The value a slot stores in its key, and the hash that places it. A short + /// word is its own key; a longer one keys on its hash and has to be confirmed + /// against [`WordCache::word_bytes_arena`], since two long words can hash alike. + /// + /// A packed key is hashed as one `u128`, not as the word's bytes it was built + /// from. Hashing a slice makes aHash mix the length in and then branch on it + /// to choose a read width; a `u128` is one fixed-width fold with nothing to + /// decide. The key already carries the length, so nothing is lost. + fn make_word_key(&self, word: &[u8]) -> (u128, u64) { + match pack_word(word) { + Some(packed) => (packed, self.hasher.hash_one(packed)), + None => { + let hash = self.hasher.hash_one(word); + ((hash as u128) | KEY_IS_HASH, hash) + } + } + } + + /// One walk over `word`'s window, from its home slot, answering both + /// questions a lookup has: which slot holds the word, and which slot it should + /// take if none of them does. + /// + /// The second answer comes out of the same walk, so a word that is not in the + /// table is placed without hashing it again or reading the tags again. + /// + /// The walk stops at the first empty slot, which is safe because storing stops + /// there too: an empty slot means the word was never stored. + fn find_word_in_cache<'w>(&self, key: u128, hash: u64, word: &'w [u8]) -> Walk<'w> { + let home = hash as usize & self.index_mask; + let tag = make_tag(hash); + // Only a hashed key can turn out to belong to a different word, and + // whether this one is hashed is settled before the walk starts, since the + // caller built the key out of the word it is looking for. + let hashed = key & KEY_IS_HASH != 0; + let mut free = None; + for step in 0..WALK_WINDOW { + let index = (home + step) & self.index_mask; + let slot_tag = self.tags[index]; + if slot_tag == EMPTY { + free = Some(index); + break; + } + if slot_tag == tag { + let slot = &self.cached_words[index]; + if slot.word_bytes_or_hash == key + && (!hashed + || self.word_bytes_arena.get(slot.word_off(), slot.word_len()) == word) + { + return Walk::Found(index); + } + } + } + Walk::Absent(Placement { + // A full window means one of its words gets evicted, and the one that + // does is the word in the home slot. The module docs say why the + // choice does not have to be any cleverer than that. + index: free.unwrap_or(home), + key, + tag, + word, + }) + } + + /// The entry to store, with whatever does not fit in a slot copied into the + /// arenas. `None` when an arena is full, so a rejected insert leaves the table + /// and the arenas exactly as they were. + fn build_entry( + &mut self, + key: u128, + word: &[u8], + ids: impl ExactSizeIterator, + ) -> Option { + let long = key & KEY_IS_HASH != 0; + let ids_len = ids.len(); + if !long && ids_len <= MAX_INLINE_IDS { + let mut payload = [0u32; MAX_INLINE_IDS]; + for (dst, id) in payload.iter_mut().zip(ids) { + *dst = id; + } + return Some(CachedWord { + word_bytes_or_hash: key, + word_ids: payload, + inline_id_count: ids_len as u8, + }); + } + // A short word stays in the key, which is a zero-length run here. + let word_len = if long { word.len() } else { 0 }; + let word_off = self.word_bytes_arena.alloc(word_len)?; + let Some(ids_off) = self.token_ids_arena.alloc(ids_len) else { + self.word_bytes_arena.release(word_off, word_len); + return None; + }; + self.word_bytes_arena + .fill(word_off, word_len, word.iter().copied()); + self.token_ids_arena.fill(ids_off, ids_len, ids); + let lengths = (word_len as u32) << PACKED_LEN_BITS | ids_len as u32; + Some(CachedWord { + word_bytes_or_hash: key, + word_ids: [word_off, ids_off, lengths], + inline_id_count: SPILLED, + }) + } + + /// Hand an overwritten entry's arena runs back, so the next entry can use + /// them. + /// + /// `#[inline]` because this exists to give a step of `insert` a name, not to + /// be called: left out of line it costs `insert` a call and the pointers it + /// had already loaded. + #[inline] + fn reclaim(&mut self, index: usize) { + let slot = self.cached_words[index]; + if !slot.ids_stored_in_arena() { + return; + } + self.word_bytes_arena + .release(slot.word_off(), slot.word_len()); + self.token_ids_arena.release(slot.ids_off(), slot.ids_len()); + } +} + +// ---------------------------------------------------------------- what a lookup returns + +/// What [`WordCache::lookup`] found. +pub enum Lookup<'c, 'w> { + /// The ids the word encoded to last time. + Hit(&'c [u32]), + /// The word is not in the table. The [`Placement`] is the slot it should go + /// in. Hand it to [`WordCache::insert`] once the model has done the work, or + /// drop it and nothing is stored. `None` when the word is too long to be + /// worth a slot at all. + Miss(Option>), +} + +impl<'c> Lookup<'c, '_> { + /// The ids, throwing the [`Placement`] away. Every caller in the encoder + /// wants the placement, so this is for tests asserting on what the table holds. + #[cfg(test)] + pub fn hit(self) -> Option<&'c [u32]> { + match self { + Lookup::Hit(ids) => Some(ids), + Lookup::Miss(_) => None, + } + } +} + +/// Where a word that missed will go, and what [`WordCache::insert`] needs to put +/// it there. Built by [`WordCache::find_word_in_cache`] out of what the walk had +/// already worked out, so storing the word costs no second hash, no second walk +/// over the tags and no second tag. +/// +/// It carries the word rather than letting `insert` take it again, because the +/// slot, the key and the tag inside were chosen for *this* word: handing back a +/// different one would file its ids under the first word's name. +pub struct Placement<'w> { + index: usize, + key: u128, + tag: u8, + word: &'w [u8], +} + +/// How a walk over a word's window ended. +enum Walk<'w> { + /// The word is in this slot. + Found(usize), + /// The word is not in the table; here is the slot it should take. + Absent(Placement<'w>), +} + +// ---------------------------------------------------------------- what a slot holds + +/// How many ids a [`CachedWord`] holds before it has to spill into +/// [`WordCache::token_ids_arena`]. Three is enough for most words in an alphabetic +/// script, and for far fewer of them in Chinese or Korean, where a word turns into +/// more ids. +const MAX_INLINE_IDS: usize = 3; + +/// A sentinel value stored in [`CachedWord::word_bytes_or_hash`] when it holds the hash of a long word instead of the word itself. +const KEY_IS_HASH: u128 = 1 << 127; + +/// A sentinel value stored in [`CachedWord::inline_id_count`] when the ids did not fit in +/// the slot and went to the arena instead. +const SPILLED: u8 = u8::MAX; + +/// Bits per length in a spilled entry's third `word_ids` lane, which packs the +/// word's byte count above the id count. Eleven, because the two have to share one +/// `u32`, and [`MAX_WORD_BYTES`] is the most either of them can be: the cache turns +/// longer words away, and no model emits more than one id per byte. +const PACKED_LEN_BITS: u32 = 11; + +const PACKED_LEN_MASK: u32 = (1 << PACKED_LEN_BITS) - 1; + +const _: () = assert!(MAX_WORD_BYTES <= PACKED_LEN_MASK as usize); + +/// One entry, in the three shapes the module docs draw out. In short: +/// +/// - `word_bytes_or_hash` is the key: the word itself when it fits in 15 bytes, +/// otherwise its hash with [`KEY_IS_HASH`] set. +/// - `word_ids` is the token ids while `inline_id_count` counts them, and becomes +/// `[word_off, ids_off, packed lengths]` once that byte is [`SPILLED`], which is +/// what [`CachedWord::word_off`] and the readers beside it are for. +/// +/// Nothing in here is read until the slot's tag has said the word might be in it. +#[derive(Clone, Copy, Default)] +#[repr(C)] +struct CachedWord { + word_bytes_or_hash: u128, + word_ids: [u32; MAX_INLINE_IDS], + inline_id_count: u8, +} + +const _: () = assert!(std::mem::size_of::() == 32); + +impl CachedWord { + fn ids_stored_in_arena(&self) -> bool { + self.inline_id_count == SPILLED + } + + fn word_off(&self) -> u32 { + self.word_ids[0] + } + + fn ids_off(&self) -> u32 { + self.word_ids[1] + } + + fn ids_len(&self) -> usize { + (self.word_ids[2] & PACKED_LEN_MASK) as usize + } + + fn word_len(&self) -> usize { + (self.word_ids[2] >> PACKED_LEN_BITS) as usize + } +} + +/// A word of `1..=15` bytes packed into a `u128`: bytes in the low lanes, length +/// in the top byte. Including the length keeps `"a"` and `"a\0"` apart, and the +/// whole key comparison becomes one register-wide equality instead of a `memcmp` +/// against bytes somewhere else in memory. +/// +/// `None` for anything longer, which is the caller's signal to key on a hash +/// instead. +/// +/// TODO: the copy is a call into `memcpy`, about a third of what building a key +/// costs, because `len` is only known at run time and LLVM folds a copy into a load +/// only when the length is a constant. A caller that knows what surrounds the word +/// could pass the 16 bytes starting where the word does instead, turning the copy +/// into one load plus a mask for the surplus bytes. +fn pack_word(word: &[u8]) -> Option { + let len = word.len(); + if len == 0 || len > 15 { + return None; + } + let mut lanes = [0u8; 16]; + lanes[..len].copy_from_slice(word); + Some(u128::from_le_bytes(lanes) | ((len as u128) << 120)) +} + +// ---------------------------------------------------------------- overflow storage + +/// A grow-only buffer that reuses the space of evicted entries instead of +/// compacting. +/// +/// Freed runs go on a free list per exact length. That is only practical because +/// [`MAX_WORD_BYTES`] bounds every run: there is a list for every length a run +/// can have, so a freed run is always reusable by the next word of the same shape +/// and no length is ever rounded up to a bigger class. The price is +/// `MAX_WORD_BYTES + 1` empty `Vec`s per arena. +struct Arena { + data: Vec, + free: Box<[Vec]>, + /// Ceiling on `data`, not a reservation. + budget: usize, +} + +impl Arena { + fn new(budget: usize) -> Self { + Self { + data: Vec::new(), + free: (0..=MAX_WORD_BYTES).map(|_| Vec::new()).collect(), + budget, + } + } + + /// A run of `len` items, or `None` once the budget is spent. + fn alloc(&mut self, len: usize) -> Option { + if let Some(off) = self.free[len].pop() { + return Some(off); + } + if self.data.len() + len > self.budget { + return None; + } + let off = self.data.len() as u32; + self.data.resize(self.data.len() + len, T::default()); + Some(off) + } + + fn release(&mut self, off: u32, len: usize) { + self.free[len].push(off); + } + + fn get(&self, off: u32, len: usize) -> &[T] { + &self.data[off as usize..off as usize + len] + } + + fn fill(&mut self, off: u32, len: usize, values: impl Iterator) { + for (dst, value) in self.data[off as usize..off as usize + len] + .iter_mut() + .zip(values) + { + *dst = value; + } + } +} + +// ---------------------------------------------------------------- the row of tags + +/// How many slots [`WordCache::find_word_in_cache`] walks before it stops looking +/// and evicts: a word is in the sixteen slots from its home slot on, or nowhere. +/// +/// Sixteen one-byte tags cross at most two cache lines, and a walk rarely reads +/// that many, since it stops at the first empty slot, usually a step or two along. +const WALK_WINDOW: usize = 16; + +/// The tag of a slot that holds nothing. No live entry ever carries it +/// ([`make_tag`]), which is what lets a walk stop at the first one it reads, +/// and lets [`WordCache::insert`] tell a slot it has to clear from one it can write +/// straight into. +const EMPTY: u8 = 0; + +/// How many bits of the hash a tag carries. All eight of them: a walk therefore +/// reads a slot it did not want about once in 255. +const TAG_BITS: u32 = 8; + +/// The tag a hash gives its slot, which is never [`EMPTY`]. +/// +/// The top of the hash, because the bottom of it already chose the home slot: a +/// tag built from those bits would be the same for every slot the word can reach +/// and would tell a walk nothing. +/// +/// Top bytes of 0 and 1 both come out as 1, since [`EMPTY`] is 0 and no live tag may +/// be that. A word with either of those top bytes therefore shares its tag with one +/// other, and pays a wasted entry read twice as often as the rest. A live tag of +/// [`EMPTY`] would cost far more: a walk would stop at that slot, losing every entry +/// stored after it in the window, and the insert that overwrote it would hand its +/// arena runs back to nobody. +fn make_tag(hash: u64) -> u8 { + ((hash >> (u64::BITS - TAG_BITS)) as u8).max(EMPTY + 1) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Look a word up and store what it encoded to, for tests that care about + /// what the table ends up holding rather than about the two steps. + fn store(cache: &mut WordCache, word: &[u8], ids: impl ExactSizeIterator) { + let at = match cache.lookup(word) { + Lookup::Miss(at) => at, + Lookup::Hit(_) => None, + }; + if let Some(at) = at { + cache.insert(at, ids); + } + } + + /// The slot a word is in, or `None` if the table does not hold it. + fn slot_of(cache: &WordCache, word: &[u8]) -> Option { + let (key, hash) = cache.make_word_key(word); + match cache.find_word_in_cache(key, hash, word) { + Walk::Found(index) => Some(index), + Walk::Absent(_) => None, + } + } + + #[test] + fn roundtrip() { + let mut cache = WordCache::new(1 << 8); + assert_eq!(cache.lookup(b"hello").hit(), None); + store(&mut cache, b"hello", [1u32, 2, 3].into_iter()); + store(&mut cache, b"world", [4u32].into_iter()); + assert_eq!(cache.lookup(b"hello").hit(), Some(&[1u32, 2, 3][..])); + assert_eq!(cache.lookup(b"world").hit(), Some(&[4u32][..])); + assert_eq!(cache.lookup(b"hell").hit(), None); + } + + /// The two properties the key encoding rests on: a packed key is unique per + /// word, and never looks like a hashed one. + #[test] + fn packed_keys_are_unique_and_never_look_hashed() { + assert_ne!(pack_word(b"a"), pack_word(b"a\0")); + assert_eq!(pack_word(&[0u8; 15]).unwrap() & KEY_IS_HASH, 0); + assert_eq!(pack_word(b""), None); + assert_eq!(pack_word(&[b'x'; 16]), None); + } + + /// A live entry's tag has to be something [`EMPTY`] is not, or a walk reads an + /// occupied slot as the end of the chain, stores on top of it, and drops the + /// arena runs it was using on the floor. Every tag is a byte of hash now, so + /// the hashes whose top byte is zero are the ones that have to be caught. + #[test] + fn a_live_tag_is_never_the_empty_marker() { + for hash in [0u64, 1, u64::MAX, 1 << 57, u64::MAX >> TAG_BITS] { + assert_ne!(make_tag(hash), EMPTY, "hash {hash:#x}"); + } + // Any other top byte is a tag in its own right, kept as it is. + assert_eq!(make_tag(0xA7 << (u64::BITS - TAG_BITS)), 0xA7); + } + + /// A short word's placement comes out of its packed key rather than its + /// bytes, which leaves the hash doing all of the mixing. Words that differ in + /// one byte pack to keys that differ in one byte, so an index taken from those + /// bits as they are (`packed as u64`) drops most of these on one slot. + #[test] + fn short_words_spread_across_the_table() { + let cache = WordCache::new(1 << 12); + let homes: std::collections::HashSet = (0..1000) + .map(|i| { + let (_, hash) = cache.make_word_key(format!("tok{i}").as_bytes()); + hash as usize & cache.index_mask + }) + .collect(); + // 1000 words over 4096 slots share homes by chance alone; ~887 distinct is + // as good as a perfect hash gets, so the floor is well under it. + assert!(homes.len() > 820, "only {} distinct homes", homes.len()); + } + + #[test] + fn oversized_words_are_ignored() { + let mut cache = WordCache::new(1 << 8); + let big = vec![7u8; MAX_WORD_BYTES + 1]; + store(&mut cache, &big, [1u32].into_iter()); + assert_eq!(cache.lookup(&big).hit(), None); + } + + /// Both sides of the 15-byte packing boundary and both sides of the inline/ + /// arena boundary have to survive a round trip, including the long word whose + /// stored `key` is only a hash and needs the byte comparison to confirm it. + #[test] + fn every_slot_shape_round_trips() { + let mut cache = WordCache::new(1 << 8); + let long = vec![b'x'; 200]; + let cases: [(&[u8], Vec); 6] = [ + (b"short", vec![1]), + (b"short-wide", (0..40).collect()), + (b"fifteen-bytes.", vec![2]), + (b"sixteen-bytes.aa", vec![3]), + (&long, vec![9]), + (&long[..64], (0..64).collect()), + ]; + for (word, ids) in &cases { + store(&mut cache, word, ids.clone().into_iter()); + } + for (word, ids) in &cases { + assert_eq!(cache.lookup(word).hit(), Some(&ids[..]), "{word:?}"); + } + } + + /// A spilled entry keeps the word's byte count and its id count in eleven bits + /// each, side by side in one `u32`. The longest word the cache accepts, when it + /// encodes to one id per byte, is the largest either of them can get. If they + /// ever overlapped, a hit would read the wrong lengths and hand back the wrong + /// ids. + #[test] + fn a_word_at_the_length_limit_round_trips() { + let mut cache = WordCache::new(1 << 8); + let word = vec![b'q'; MAX_WORD_BYTES]; + let ids: Vec = (0..MAX_WORD_BYTES as u32).collect(); + store(&mut cache, &word, ids.clone().into_iter()); + assert_eq!(cache.lookup(&word).hit(), Some(&ids[..])); + } + + /// A tag is one byte, so one slot in 255 answers for a word that is not in + /// it. That is too rare to wait for, so forge one: put a word's tag on a slot + /// holding a different key, and demand the walk look past it rather than + /// treat the tag as the answer. + #[test] + fn a_tag_collision_is_confirmed_against_the_key() { + let mut cache = WordCache::new(WALK_WINDOW); + let (_, hash) = cache.make_word_key(b"beta"); + let decoy = hash as usize & cache.index_mask; + cache.tags[decoy] = make_tag(hash); + // Any key that is not beta's. A packed key always carries a length in its + // top byte, so 1 cannot be one. + cache.cached_words[decoy] = CachedWord { + word_bytes_or_hash: 1, + word_ids: [7, 0, 0], + inline_id_count: 1, + }; + + store(&mut cache, b"beta", [2u32].into_iter()); + assert_eq!(cache.lookup(b"beta").hit(), Some(&[2u32][..])); + } + + /// Two long words can hash to the same key, and then only their bytes tell + /// them apart. Real hash collisions are too rare to write a test around, so + /// forge one: park another word's entry on this word's home slot, stamp this + /// word's key and tag on it, and demand the walk look past it. + #[test] + fn a_hashed_key_is_confirmed_against_the_word_bytes() { + let mut cache = WordCache::new(1 << 8); + let mine = vec![b'a'; 40]; + let theirs = vec![b'b'; 40]; + + store(&mut cache, &theirs, [7u32].into_iter()); + let their_index = slot_of(&cache, &theirs).unwrap(); + let their_slot = cache.cached_words[their_index]; + let (my_key, my_hash) = cache.make_word_key(&mine); + let my_home = my_hash as usize & cache.index_mask; + cache.cached_words[my_home] = CachedWord { + word_bytes_or_hash: my_key, + ..their_slot + }; + cache.tags[my_home] = make_tag(my_hash); + + store(&mut cache, &mine, [1u32, 2].into_iter()); + assert_eq!(cache.lookup(&mine).hit(), Some(&[1u32, 2][..])); + } + + /// A word that hashes into a full window is stored rather than turned away, + /// and the entry evicted is the one in its home slot, even when that entry has + /// just been used, which is the whole of what this policy costs. + #[test] + fn a_full_window_evicts_its_home_slot() { + let mut cache = WordCache::new(WALK_WINDOW); + let words: Vec> = (0..WALK_WINDOW as u8).map(|i| vec![i; 4]).collect(); + for (i, word) in words.iter().enumerate() { + store(&mut cache, word, [i as u32].into_iter()); + } + let (_, hash) = cache.make_word_key(b"newcomer"); + let home = hash as usize & cache.index_mask; + let evicted = words + .iter() + .position(|word| slot_of(&cache, word) == Some(home)) + .expect("the table is full, so some word holds that slot"); + for _ in 0..8 { + assert!(cache.lookup(&words[evicted]).hit().is_some()); + } + + store(&mut cache, b"newcomer", [999u32].into_iter()); + assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999u32][..])); + assert_eq!(cache.lookup(&words[evicted]).hit(), None); + for (i, word) in words.iter().enumerate() { + if i != evicted { + assert_eq!( + cache.lookup(word).hit(), + Some(&[i as u32][..]), + "entry {i} was dropped as well" + ); + } + } + } + + /// A window that runs off the end of the table carries on at the start, and a + /// word stored past that seam has to be found there. Nothing mirrors the first + /// tags at the end of the row. Every step of the walk is folded back into the + /// table instead. + #[test] + fn a_word_stored_past_the_end_of_the_table_is_found() { + let mut cache = WordCache::new(WALK_WINDOW); + let last_slot = cache.cached_words.len() - 1; + let homed_on_the_last_slot: Vec> = (0..4000u32) + .map(|i| format!("w{i}").into_bytes()) + .filter(|word| { + let (_, hash) = cache.make_word_key(word); + hash as usize & cache.index_mask == last_slot + }) + .take(3) + .collect(); + assert_eq!(homed_on_the_last_slot.len(), 3); + + for (i, word) in homed_on_the_last_slot.iter().enumerate() { + store(&mut cache, word, [i as u32].into_iter()); + } + for (i, word) in homed_on_the_last_slot.iter().enumerate() { + assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); + } + // The first one took the last slot, so the other two could only go past it. + assert!(slot_of(&cache, &homed_on_the_last_slot[1]).unwrap() < last_slot); + } + + /// Reusing an evicted entry's arena run is where this design can go wrong: + /// hand one run to two live entries and a hit starts returning another word's + /// ids. Churn a table far too small for the input and demand the invariant + /// that matters: an entry may be evicted, but a hit is never wrong. + #[test] + fn a_hit_never_returns_another_words_ids() { + let mut cache = WordCache::new(64); + let mut expected: Vec<(Vec, Vec)> = Vec::new(); + for i in 0..2000usize { + let word = match i % 4 { + 0 => format!("w{i}"), + 1 => format!("a-long-word-past-fifteen-bytes-{i}"), + 2 => format!("k{i}xxxxxxxxxxxx"), + _ => format!("{}-{i}", "z".repeat(i % 40)), + } + .into_bytes(); + let ids: Vec = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); + store(&mut cache, &word, ids.clone().into_iter()); + expected.push((word, ids)); + } + let mut live = 0; + for (word, ids) in &expected { + if let Some(hit) = cache.lookup(word).hit() { + assert_eq!(hit, &ids[..], "{word:?}"); + live += 1; + } + } + assert!( + live > 0, + "everything was evicted, so the test proves nothing" + ); + } + + /// Freed runs have to come back. Without reuse the arenas grow with every + /// insert until their budget is spent, and from then on the cache silently + /// stops accepting words even though the table has room for them. + #[test] + fn arenas_hold_the_live_set_not_every_insert() { + let mut cache = WordCache::new(64); + // Same length every time, so one free list serves every word and the + // bounds below are exact: 64 slots, plus the run reserved for the insert + // in flight before the entry it replaces gives its own run back. + let word = |i: usize| format!("a-long-word-past-fifteen-bytes-{i:04}"); + for i in 0..5000usize { + store(&mut cache, word(i).as_bytes(), [i as u32; 8].into_iter()); + } + assert_eq!( + cache.lookup(word(4999).as_bytes()).hit(), + Some(&[4999u32; 8][..]), + "inserts stopped landing: the arenas ran out" + ); + assert!(cache.word_bytes_arena.data.len() <= 65 * 35); + assert!(cache.token_ids_arena.data.len() <= 65 * 8); + } +} From a2bde6417c5c1833c0f8bcaebb5754452dfc1822 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:37:31 +0200 Subject: [PATCH 02/13] perf(atomsplit): SIMD literal pattern search (#2271) A vectorized literal scanner in atomsplit, wired into the Replace normalizer and the Split pre-tokenizer so literal patterns stop going through the regex engine. Squashed from feat/literal-simd. --- tokenizers/atomsplit/Cargo.toml | 4 + tokenizers/atomsplit/benches/literal.rs | 67 ++ tokenizers/atomsplit/src/lib.rs | 6 + tokenizers/atomsplit/src/literal.rs | 141 +++- tokenizers/atomsplit/src/simd_literal.rs | 783 ++++++++++++++++++ tokenizers/atomsplit/tests/literal.rs | 200 +++++ .../tk-encode/src/normalizers/replace.rs | 20 +- .../tk-encode/src/pre_tokenizers/split.rs | 106 ++- tokenizers/tk-encode/src/tokenizer/pattern.rs | 13 +- .../tk-encode/src/tokenizer/pipeline.rs | 154 ++-- 10 files changed, 1417 insertions(+), 77 deletions(-) create mode 100644 tokenizers/atomsplit/benches/literal.rs create mode 100644 tokenizers/atomsplit/src/simd_literal.rs diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml index fb8947f3f4..dce6648808 100644 --- a/tokenizers/atomsplit/Cargo.toml +++ b/tokenizers/atomsplit/Cargo.toml @@ -43,6 +43,10 @@ harness = false name = "classify" harness = false +[[bench]] +name = "literal" +harness = false + [[bench]] name = "class_runs" harness = false diff --git a/tokenizers/atomsplit/benches/literal.rs b/tokenizers/atomsplit/benches/literal.rs new file mode 100644 index 0000000000..c165b3f6c2 --- /dev/null +++ b/tokenizers/atomsplit/benches/literal.rs @@ -0,0 +1,67 @@ +//! `Literal::matches` (one search per match) against `Literal::matches_into` (one scan per +//! text), across the delimiter densities pre-tokenizers see: a space about every six bytes of +//! English, a `▁` per word after a SentencePiece replace, and a pattern the text does not +//! contain at all (where `memmem`'s skip-ahead is at its best and the scan must keep up). +//! Run: cargo bench --bench literal +use atomsplit::literal::Literal; +use std::hint::black_box; +use std::time::Instant; + +fn best_ns_per_byte(text_len: usize, mut pass: impl FnMut() -> usize) -> f64 { + let iters = (16_000_000 / text_len.max(1)).clamp(4, 400) as u32; + for _ in 0..3 { + black_box(pass()); + } + let mut best = f64::INFINITY; + for _ in 0..9 { + let t = Instant::now(); + for _ in 0..iters { + black_box(pass()); + } + best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text_len) as f64); + } + best +} + +fn main() { + let manifest = env!("CARGO_MANIFEST_DIR"); + let english = + std::fs::read_to_string(format!("{manifest}/../data/big.txt")).unwrap_or_default(); + let english: String = english.chars().take(2_000_000).collect(); + if english.is_empty() { + println!("empty corpus — data/big.txt missing? (make test downloads it)"); + return; + } + let metaspaced = english.replace(' ', "\u{2581}"); + + println!( + "{:<22} {:>9} {:>12} {:>12} {:>9} {:>12}", + "", "matches", "iterator", "batch", "speedup", "count-only" + ); + for (label, text, pattern) in [ + ("space in English", english.as_bytes(), " "), + ("▁ in metaspaced", metaspaced.as_bytes(), "\u{2581}"), + ("▁ absent", english.as_bytes(), "\u{2581}"), + ] { + let literal = Literal::new(pattern.as_bytes()).unwrap(); + let mut offsets: Vec = Vec::with_capacity(text.len()); + let mut buffer = vec![0u32; text.len() + 4]; + + let count = literal.matches(text).count(); + let iterator = best_ns_per_byte(text.len(), || { + offsets.clear(); + offsets.extend(literal.matches(text)); + offsets.len() + }); + let batch = best_ns_per_byte(text.len(), || literal.matches_into(text, &mut buffer)); + let counting = best_ns_per_byte(text.len(), || literal.count_matches(text)); + + println!( + "{label:<22} {count:>9} {:>7.2} GB/s {:>7.2} GB/s {:>8.2}x {:>7.2} GB/s", + 1.0 / iterator, + 1.0 / batch, + iterator / batch, + 1.0 / counting + ); + } +} diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs index 5f22f52d00..13beac00af 100644 --- a/tokenizers/atomsplit/src/lib.rs +++ b/tokenizers/atomsplit/src/lib.rs @@ -28,6 +28,12 @@ mod simd_avx_classify; #[cfg(target_arch = "aarch64")] mod simd_classify; mod simd_fsm; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod simd_literal; #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] mod simd_wasm_classify; pub mod tables; diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs index c826b2dda6..9e7106c1ad 100644 --- a/tokenizers/atomsplit/src/literal.rs +++ b/tokenizers/atomsplit/src/literal.rs @@ -4,7 +4,15 @@ //! FSMs cut where the class changes ([`crate::fsm`]). //! //! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string: -//! We can use a simpler byte search looking at 16 bytes at a time. +//! a plain byte search finds the same cuts. Two searches cover the two match densities: +//! +//! - [`Literal::matches`] iterates with [`memmem`], which is built for needles that are rare in +//! the haystack: it skips ahead fast and restarts after every match. +//! - [`Literal::matches_into`] scans the whole text once and writes every offset into a caller +//! buffer. Pre-tokenizer delimiters are the opposite of rare (running English text has a space +//! about every six bytes), and at that density the per-match restarts dominate the iterator. +//! The scan answers a whole block of text at once and turns the answers into offsets; +//! `simd_literal.rs` explains how, step by step. use memchr::memmem; use std::fmt; @@ -53,4 +61,135 @@ impl Literal { pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator + 't { self.finder.find_iter(text) } + + /// The same offsets as [`Literal::matches`], written into `out[..count]` in one scan of + /// `text`; returns the count. Use this when matches are frequent: the iterator restarts its + /// search machinery after every match, the scan never stops. + /// + /// # Preconditions + /// - `out.len() >= text.len() / pattern.len() + 4` (asserted). The division is the most + /// matches the text can hold; the `+ 4` is slack the SIMD path writes past the last match. + /// - `text.len()` must fit in `u32` (asserted), so that every offset does too. + pub fn matches_into(&self, text: &[u8], out: &mut [u32]) -> usize { + assert!( + out.len() >= text.len() / self.pattern().len() + 4, + "matches_into needs out.len() >= text.len() / pattern.len() + 4" + ); + assert!( + u32::try_from(text.len()).is_ok(), + "matches_into writes u32 offsets" + ); + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + if self.scannable() { + match *self.pattern() { + [a] => return crate::simd_literal::matches_into([a], text, out), + [a, b] => return crate::simd_literal::matches_into([a, b], text, out), + [a, b, c] => return crate::simd_literal::matches_into([a, b, c], text, out), + _ => unreachable!("scannable patterns are one to three bytes"), + } + } + let mut count = 0; + for position in self.finder.find_iter(text) { + out[count] = position as u32; + count += 1; + } + count + } + + /// How many matches [`Literal::matches`] would report. Counting runs the compare step + /// without producing any offsets, so it is several times faster than a full scan; count + /// first to size an output exactly, then fill it with a [`Literal::for_each_match`] pass. + #[must_use] + pub fn count_matches(&self, text: &[u8]) -> usize { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + if self.scannable() { + match *self.pattern() { + [a] => return crate::simd_literal::count_matches([a], text), + [a, b] => return crate::simd_literal::count_matches([a, b], text), + [a, b, c] => return crate::simd_literal::count_matches([a, b, c], text), + _ => unreachable!("scannable patterns are one to three bytes"), + } + } + self.finder.find_iter(text).count() + } + + /// Calls `on_match` with the byte offset of every match, left to right: the same offsets + /// as [`Literal::matches`] at batch-scan speed, with no buffer for the caller to provide. + /// The scan streams through a small stack window, so its footprint stays flat however + /// long the text is. + pub fn for_each_match(&self, text: &[u8], mut on_match: impl FnMut(usize)) { + if !self.scannable() { + for position in self.finder.find_iter(text) { + on_match(position); + } + return; + } + self.scan_windows(text, |base, matches| { + for &position in matches { + on_match(base + position as usize); + } + }); + } + + /// Whether the batch scan covers this pattern. The scan emits every position where the + /// pattern matches, with no ordering check between them; that is only the non-overlapping + /// match list when the pattern cannot overlap itself, which one byte comparison per + /// length decides: a two-byte pattern overlaps itself when both bytes are equal, a + /// three-byte one when last equals first. Longer or self-overlapping patterns search + /// through the [`memmem`] engine instead. + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + fn scannable(&self) -> bool { + match *self.pattern() { + [_] => true, + [a, b] => a != b, + [a, _, c] => a != c, + _ => false, + } + } + + /// Without a SIMD kernel there is no batch scan; everything searches through [`memmem`]. + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + fn scannable(&self) -> bool { + false + } + + /// Streams the batch scan window by window: [`Literal::matches_into`] fills a stack + /// buffer for one stretch of text, `on_window` consumes it (offsets are relative to the + /// window's `base`), and the next window starts where a match could no longer fit in the + /// previous one, so nothing at a window edge is missed or reported twice. + fn scan_windows(&self, text: &[u8], mut on_window: impl FnMut(usize, &[u32])) { + // 4KB of stack; a window of (1024 - 4) * pattern.len() bytes fills it exactly. + let mut buffer = [0u32; 1024]; + let width = self.pattern().len(); + let window = (buffer.len() - 4) * width; + let mut base = 0; + loop { + let end = usize::min(base + window, text.len()); + let count = self.matches_into(&text[base..end], &mut buffer); + on_window(base, &buffer[..count]); + if end == text.len() { + return; + } + // A match crossing `end` was not reported; the next window re-covers its + // possible starts. Matches of a scannable pattern never overlap, so the reports + // stay disjoint. + base = end - (width - 1); + } + } } diff --git a/tokenizers/atomsplit/src/simd_literal.rs b/tokenizers/atomsplit/src/simd_literal.rs new file mode 100644 index 0000000000..78e6bfa587 --- /dev/null +++ b/tokenizers/atomsplit/src/simd_literal.rs @@ -0,0 +1,783 @@ +//! Finds every place a short pattern (up to three bytes) occurs in a text, in one pass. +//! This is the scan behind [`Literal::matches_into`](crate::literal::Literal::matches_into); +//! [`crate::literal`] explains when it is picked over the plain iterator. +//! +//! # Why a dedicated scan +//! +//! A pre-tokenizer splits on a delimiter that shows up every few bytes: running English text +//! has a space about every six bytes, SentencePiece text has a `▁` at every word. A +//! next-match engine like `memmem` is built for the opposite case, a needle that is rare, so +//! at this density it spends most of its time stopping and restarting: every match returns to +//! the caller, and every call re-enters the search machinery. This scan never stops. It +//! answers "where are ALL the matches in these 16 bytes?" with a handful of instructions, +//! block after block. +//! +//! # How: compare, mask, decode +//! +//! Two words come up in everything below. A SIMD **register** is an extra-wide CPU value +//! holding 16 bytes side by side; one instruction operates on all 16 at once. Each of those +//! 16 byte slots is a **lane**. A SIMD compare answers lane by lane: a lane becomes all ones +//! where the answer is yes, all zeros where it is no. +//! +//! **Step 1, compare.** Compare the text against the pattern's first byte (copied across all +//! 16 lanes), the text shifted by one against the second byte, and so on, then AND the +//! answers together: a lane that passed every compare is the start of a whole match, and +//! nothing needs a second look. Searching for `▁` (three bytes, `E2 96 81`) in +//! `"a…b▁c"`, where the `…` (`E2 80 A6`) shares a first byte with `▁` and acts as the decoy: +//! +//! ```text +//! offset 0 1 2 3 4 5 6 7 8 +//! text a E2 80 A6 b E2 96 81 c +//! └── … ──────┘ └── ▁ ──────┘ +//! text[j] == E2 ? · ✓ · · · ✓ · · · +//! text[j+1] == 96 ? · · · · · ✓ · · · +//! text[j+2] == 81 ? · · · · · ✓ · · · +//! AND of the three · · · · · ✓ · · · +//! ``` +//! +//! Only offset 5 passes all three rows: a match starts at 5, and the decoy at 1 is out after +//! the second row. +//! +//! **Step 2, mask.** The 16 lane answers are squeezed into one ordinary integer so plain +//! arithmetic can take over: bit `j` set means "a match starts at offset `j`". For the +//! example above the mask is `0b100000`, bit 5. Each target has a one-or-two instruction way +//! to build this integer; on NEON the bits come out spaced four positions apart instead of +//! one, a spacing the shared code carries along as `SHIFT` (the per-target layer below). +//! +//! **Step 3, decode.** Count-trailing-zeros on the mask gives the lowest set bit, which is the +//! next match offset; clear that bit and repeat. The one twist is that the code decodes four +//! offsets *without checking whether any bits are left*. Say a block starting at text offset +//! `base` had matches at offsets 2 and 5 (shown with the one-bit spacing): +//! +//! ```text +//! mask 0b100100 trailing zeros = 2 write base + 2 clear bit → 0b100000 +//! mask 0b100000 trailing zeros = 5 write base + 5 clear bit → 0 +//! mask 0 trailing zeros = 64 write garbage (slot 3) +//! mask 0 trailing zeros = 64 write garbage (slot 4) +//! cursor += 2 (how many bits were set) +//! ``` +//! +//! The two garbage slots sit past the cursor, so the next block's writes (or nothing) land on +//! them and they never reach the caller. The honest alternative, checking "still something +//! left?" before each decode, puts a branch in the hottest loop; with a match every four to +//! eight bytes that branch guesses wrong about once per block, which costs more than all the +//! compares together. +//! +//! Sparse text takes none of this: one instruction per 64-byte block answers "nothing here" +//! (see `any`), so a pattern that is rare or absent stays close to `memmem` speed. +//! +//! The compare trick is Wojciech Muła's ("SIMD-friendly algorithms for substring searching"); +//! the branch-free decode is how simdjson reads its masks. +//! +//! # The per-target layer +//! +//! Everything from the mask on is plain integer arithmetic and is shared. A target provides +//! four primitives in its `arch` module: +//! +//! - `splat`: one byte copied into every lane, the shape a compare wants. +//! - `match_starts`: step 1, the shifted compares ANDed together. +//! - `any`: "did anything in these 64 bytes match?", the cheap exit for sparse text. +//! - `bits`: step 2, lane answers to integer mask. +//! +//! x86 (`movemask`) and wasm (`bitmask`) each have an instruction that grabs one bit from +//! every lane, so offset `j` lands on bit `j`. NEON has no such instruction; its standard +//! substitute (Danila Kutenin's) leaves the bits four positions apart, so offset `j` lands +//! on bit `4 * j`. `SHIFT` is that spacing (`bit position >> SHIFT` recovers the offset), +//! and it is the only difference the shared code ever sees. +//! +//! On x86 the 16-byte baseline is SSE2, which every x86_64 CPU has; when the CPU reports +//! wider vectors, [`matches_into`] and [`count_matches`] pick a faster front end per call +//! (the `wide` module: AVX2 and AVX-512). +//! +//! The wasm layer is compile-checked but not run in CI; the other targets run the full +//! `tests/literal.rs` parity suite. + +#[cfg(target_arch = "aarch64")] +mod arch { + use core::arch::aarch64::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = uint8x16_t; + /// Offsets sit four bit positions apart in a mask from [`bits`]: offset `j` is bit `4 * j`. + pub(super) const SHIFT: u32 = 2; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + unsafe { vdupq_n_u8(byte) } + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { vceqq_u8(vld1q_u8(p), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { vandq_u8(starts, vceqq_u8(vld1q_u8(p.add(k)), byte)) }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + unsafe { vmaxvq_u8(vorrq_u8(vorrq_u8(s0, s1), vorrq_u8(s2, s3))) != 0 } + } + + /// The lane answers as one integer: one bit per lane, four bit positions apart. + /// + /// NEON cannot grab one bit from every lane in a single instruction. Instead `vshrn` + /// ("shift right and narrow") halves the register to 64 bits, shrinking each lane's + /// answer from eight bits to four (still all ones or all zeros); the `0x1111..` mask + /// then keeps one bit of each four, so clearing one bit in the decode drops exactly + /// one match. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + unsafe { + let halved = vshrn_n_u16::<4>(vreinterpretq_u16_u8(starts)); + vget_lane_u64::<0>(vreinterpret_u64_u8(halved)) & 0x1111_1111_1111_1111 + } + } +} + +#[cfg(target_arch = "x86_64")] +mod arch { + use core::arch::x86_64::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = __m128i; + /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift. + pub(super) const SHIFT: u32 = 0; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + unsafe { _mm_set1_epi8(byte as i8) } + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { _mm_cmpeq_epi8(_mm_loadu_si128(p.cast()), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { + _mm_and_si128( + starts, + _mm_cmpeq_epi8(_mm_loadu_si128(p.add(k).cast()), byte), + ) + }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + unsafe { _mm_movemask_epi8(_mm_or_si128(_mm_or_si128(s0, s1), _mm_or_si128(s2, s3))) != 0 } + } + + /// The lane answers as one integer: `movemask` grabs each lane's top bit, so offset `j` + /// is bit `j`. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + unsafe { _mm_movemask_epi8(starts) as u32 as u64 } + } +} + +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +mod arch { + use core::arch::wasm32::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = v128; + /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift. + pub(super) const SHIFT: u32 = 0; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + u8x16_splat(byte) + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { u8x16_eq(v128_load(p.cast()), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { v128_and(starts, u8x16_eq(v128_load(p.add(k).cast()), byte)) }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + v128_any_true(v128_or(v128_or(s0, s1), v128_or(s2, s3))) + } + + /// The lane answers as one integer: `bitmask` grabs each lane's top bit, so offset `j` + /// is bit `j`. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + u8x16_bitmask(starts) as u64 + } +} + +/// Decodes every set bit of `bits` into a `base + offset` value at `cursor`, and returns the +/// cursor advanced by the true match count. +/// +/// The first four decodes run without checking `bits` (the module docs say why): once the mask +/// is used up, `trailing_zeros` is 64 and the slot gets a garbage offset, scratch that the +/// count never covers and a later call overwrites. Blocks with more than four matches finish +/// in the loop. +/// +/// # Safety +/// There must be room for `max(4, count_ones(bits))` writes between `cursor` and `end`: four +/// slots for the unconditional decodes, one per match when a block has more. Debug builds +/// check this bound before writing anything. +#[inline(always)] +unsafe fn decode(mut bits: u64, base: u32, cursor: *mut u32, end: *const u32) -> *mut u32 { + let count = bits.count_ones() as usize; + debug_assert!( + end as usize - cursor as usize >= count.max(4) * size_of::(), + "decode would write past the match buffer" + ); + // SAFETY: the caller's contract above gives these four slots. + for slot in 0..4 { + unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) }; + bits &= bits.wrapping_sub(1); + } + if bits != 0 { + let mut slot = 4; + // SAFETY: one slot per remaining match, still within the caller's `count` slots. + while bits != 0 { + unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) }; + bits &= bits.wrapping_sub(1); + slot += 1; + } + } + // SAFETY: `count` slots were just written, so one past the last of them is in bounds. + unsafe { cursor.add(count) } +} + +/// How many matches of `pattern` are in `text`: the compare and mask steps only, counting set +/// bits instead of decoding them, which makes counting several times faster than the full +/// scan. Same pattern rules as [`matches_into`]. +pub(crate) fn count_matches(pattern: [u8; K], text: &[u8]) -> usize { + const { + assert!( + K >= 1 && K <= 3, + "the scan is built for patterns of 1 to 3 bytes" + ); + } + let len = text.len(); + if len < K { + return 0; + } + #[cfg(target_arch = "x86_64")] + if let Some(count) = wide::count_matches_wide::(pattern, text) { + return count; + } + let p = text.as_ptr(); + let mut pattern_chunks = [arch::splat(0); K]; + for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) { + *chunk = arch::splat(byte); + } + let mut count = 0usize; + let mut i = 0usize; + + // The same block layout as `scan`, and the same SAFETY story for the reads: each loop + // condition bounds its loads inside `text`. Nothing here writes. + while i + 64 + (K - 1) <= len { + let (s0, s1, s2, s3) = unsafe { + ( + arch::match_starts::(p.add(i), &pattern_chunks), + arch::match_starts::(p.add(i + 16), &pattern_chunks), + arch::match_starts::(p.add(i + 32), &pattern_chunks), + arch::match_starts::(p.add(i + 48), &pattern_chunks), + ) + }; + if arch::any(s0, s1, s2, s3) { + count += (arch::bits(s0).count_ones() + + arch::bits(s1).count_ones() + + arch::bits(s2).count_ones() + + arch::bits(s3).count_ones()) as usize; + } + i += 64; + } + + count + unsafe { count_tail::(pattern, &pattern_chunks, text, i) } +} + +/// The counting twin of [`scan_tail`]: 16-byte rounds from `i`, the overlapped final block, +/// or the byte-by-byte walk for texts shorter than one block. +/// +/// # Safety +/// Reads are bounded by the loop conditions, exactly as in [`scan_tail`]. Nothing writes. +unsafe fn count_tail( + pattern: [u8; K], + pattern_chunks: &[arch::Chunk; K], + text: &[u8], + mut i: usize, +) -> usize { + let len = text.len(); + let p = text.as_ptr(); + let mut count = 0usize; + + // SAFETY: reads end at `i + 16 + (K - 1) <= len`. + while i + 16 + (K - 1) <= len { + let starts = unsafe { arch::match_starts::(p.add(i), pattern_chunks) }; + count += arch::bits(starts).count_ones() as usize; + i += 16; + } + + if i + K <= len { + if len >= 16 + (K - 1) { + // The overlapped final block: bits below `i` were counted above and are masked off. + // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`. + let base = len - 16 - (K - 1); + let starts = unsafe { arch::match_starts::(p.add(base), pattern_chunks) }; + let bits = arch::bits(starts) & (!0u64 << ((i - base) << arch::SHIFT)); + count += bits.count_ones() as usize; + } else { + while i + K <= len { + count += (text[i..i + K] == pattern[..]) as usize; + i += 1; + } + } + } + count +} + +/// The wider x86 front ends. Every x86_64 CPU has the SSE2 baseline above; these two rungs +/// are picked per call when the CPU reports the feature: +/// +/// - **AVX2** compares 32 bytes per instruction, so one round covers 64 bytes with two +/// compares and one combined mask. +/// - **AVX-512** compares 64 bytes straight into a bit mask (`_mm512_cmpeq_epi8_mask`, no +/// separate mask step) and decodes with `vpcompressd`: the CPU packs the matching offsets +/// itself, 16 at a time, replacing the whole count-trailing-zeros loop. +/// +/// Each rung is its own function because `#[target_feature]` applies per function; the bodies +/// mirror [`scan`] and [`count_matches`] round for round and share [`decode`], [`scan_tail`] +/// and [`count_tail`]. The AVX2 rung is exercised by the `tests/literal.rs` suite on any +/// AVX2 machine or emulator; the AVX-512 rung runs the same suite only on hardware that has +/// it, so run the tests on such a machine before trusting changes to it. +#[cfg(target_arch = "x86_64")] +mod wide { + use super::{arch, count_tail, decode, scan_tail}; + use core::arch::x86_64::*; + + /// Runs the widest scan this CPU supports; `None` means only the baseline exists. + #[inline] + pub(super) fn matches_into_wide( + pattern: [u8; K], + text: &[u8], + out: &mut [u32], + ) -> Option { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: the features were just detected; the buffer contract is the caller's. + return Some(unsafe { scan_avx512::(pattern, text, out) }); + } + if is_x86_feature_detected!("avx2") { + // SAFETY: the feature was just detected; the buffer contract is the caller's. + return Some(unsafe { scan_avx2::(pattern, text, out) }); + } + None + } + + /// The counting twin of [`matches_into_wide`]. + #[inline] + pub(super) fn count_matches_wide( + pattern: [u8; K], + text: &[u8], + ) -> Option { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: the features were just detected. Counting never writes. + return Some(unsafe { count_avx512::(pattern, text) }); + } + if is_x86_feature_detected!("avx2") { + // SAFETY: the feature was just detected. Counting never writes. + return Some(unsafe { count_avx2::(pattern, text) }); + } + None + } + + /// All-ones per lane where a whole `K`-byte match starts, for 32 positions at `p`. + /// + /// # Safety + /// AVX2 must be available, and `p .. p + 32 + K - 1` must be readable. + #[inline] + #[target_feature(enable = "avx2")] + unsafe fn match_starts32(p: *const u8, pattern: &[__m256i; K]) -> __m256i { + // SAFETY: the caller's contract above; each load of 32 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.cast()) }, pattern[0]); + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = _mm256_and_si256( + starts, + _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.add(k).cast()) }, byte), + ); + } + starts + } + + /// One byte copied into every lane, and the baseline chunks [`scan_tail`] wants. + #[target_feature(enable = "avx2")] + fn splat32(pattern: [u8; K]) -> ([__m256i; K], [arch::Chunk; K]) { + let mut wide = [_mm256_set1_epi8(0); K]; + let mut baseline = [arch::splat(0); K]; + for k in 0..K { + wide[k] = _mm256_set1_epi8(pattern[k] as i8); + baseline[k] = arch::splat(pattern[k]); + } + (wide, baseline) + } + + /// [`scan`](super::scan) with the AVX2 front end: 64 bytes per round as two 32-byte + /// compares, one combined mask, decoded in the same 16-bit quarters as the baseline so + /// [`decode`]'s four-slot unroll stays matched to typical match density. + /// + /// # Safety + /// AVX2 must be available, and `out` must satisfy + /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts. + #[target_feature(enable = "avx2")] + unsafe fn scan_avx2(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat32::(pattern); + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read). + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + + // SAFETY: reads end at `i + 32 + 32 + (K - 1) <= len`; writes are the buffer + // contract, handed to `decode` as `end`. + while i + 64 + (K - 1) <= len { + let (s0, s1) = unsafe { + ( + match_starts32::(p.add(i), &pattern_chunks), + match_starts32::(p.add(i + 32), &pattern_chunks), + ) + }; + let bits = (_mm256_movemask_epi8(s0) as u32 as u64) + | ((_mm256_movemask_epi8(s1) as u32 as u64) << 32); + if bits != 0 { + unsafe { + cursor = decode(bits & 0xFFFF, i as u32, cursor, end); + cursor = decode((bits >> 16) & 0xFFFF, i as u32 + 16, cursor, end); + cursor = decode((bits >> 32) & 0xFFFF, i as u32 + 32, cursor, end); + cursor = decode(bits >> 48, i as u32 + 48, cursor, end); + } + } + i += 64; + } + // SAFETY: `scan_tail` continues under the same contract. + cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } + } + + /// The counting twin of [`scan_avx2`]. + /// + /// # Safety + /// AVX2 must be available. + #[target_feature(enable = "avx2")] + unsafe fn count_avx2(pattern: [u8; K], text: &[u8]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat32::(pattern); + let mut count = 0usize; + let mut i = 0usize; + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes. + while i + 64 + (K - 1) <= len { + let (s0, s1) = unsafe { + ( + match_starts32::(p.add(i), &pattern_chunks), + match_starts32::(p.add(i + 32), &pattern_chunks), + ) + }; + count += (_mm256_movemask_epi8(s0).count_ones() + _mm256_movemask_epi8(s1).count_ones()) + as usize; + i += 64; + } + // SAFETY: `count_tail` continues under the same read bounds. + count + unsafe { count_tail::(pattern, &baseline, text, i) } + } + + /// The 64-byte match mask, straight from the compares: AVX-512 byte compares produce a + /// bit per lane natively, so there is no separate mask step to pay for. + /// + /// # Safety + /// AVX-512 F and BW must be available, and `p .. p + 64 + K - 1` must be readable. + #[inline] + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn match_bits64(p: *const u8, pattern: &[__m512i; K]) -> u64 { + // SAFETY: the caller's contract above; each load of 64 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut bits = _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.cast()) }, pattern[0]); + for (k, &byte) in pattern.iter().enumerate().skip(1) { + bits &= _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.add(k).cast()) }, byte); + } + bits + } + + /// One byte copied into every lane, and the baseline chunks the shared tails want. + #[target_feature(enable = "avx512f")] + fn splat64(pattern: [u8; K]) -> ([__m512i; K], [arch::Chunk; K]) { + let mut wide = [_mm512_set1_epi8(0); K]; + let mut baseline = [arch::splat(0); K]; + for k in 0..K { + wide[k] = _mm512_set1_epi8(pattern[k] as i8); + baseline[k] = arch::splat(pattern[k]); + } + (wide, baseline) + } + + /// [`scan`](super::scan) with the AVX-512 front end. The decode is `vpcompressd` + /// (`_mm512_mask_compressstoreu_epi32`): give the CPU sixteen candidate offsets and a + /// mask, and it stores exactly the matching ones, packed. The count-trailing-zeros loop + /// and its density concerns disappear; four compress-stores cover a 64-byte round. + /// + /// # Safety + /// AVX-512 F and BW must be available, and `out` must satisfy + /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts. + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn scan_avx512(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat64::(pattern); + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read). + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + // Lane numbers 0..16, the base of every offset vector. + let lanes = _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Each compress-store writes exactly + // the matches of its quarter, within the buffer contract. + while i + 64 + (K - 1) <= len { + let bits = unsafe { match_bits64::(p.add(i), &pattern_chunks) }; + if bits != 0 { + let base = _mm512_add_epi32(lanes, _mm512_set1_epi32(i as i32)); + for quarter in 0..4 { + let mask = ((bits >> (16 * quarter)) & 0xFFFF) as __mmask16; + let offsets = _mm512_add_epi32(base, _mm512_set1_epi32(16 * quarter)); + debug_assert!( + end as usize - cursor as usize + >= mask.count_ones() as usize * size_of::(), + "compress-store would write past the match buffer" + ); + unsafe { + _mm512_mask_compressstoreu_epi32(cursor.cast(), mask, offsets); + cursor = cursor.add(mask.count_ones() as usize); + } + } + } + i += 64; + } + // SAFETY: `scan_tail` continues under the same contract. + cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } + } + + /// The counting twin of [`scan_avx512`]. + /// + /// # Safety + /// AVX-512 F and BW must be available. + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn count_avx512(pattern: [u8; K], text: &[u8]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat64::(pattern); + let mut count = 0usize; + let mut i = 0usize; + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes. + while i + 64 + (K - 1) <= len { + count += unsafe { match_bits64::(p.add(i), &pattern_chunks) }.count_ones() as usize; + i += 64; + } + // SAFETY: `count_tail` continues under the same read bounds. + count + unsafe { count_tail::(pattern, &baseline, text, i) } + } +} + +/// One-pass scan: writes the start offset of every match of `pattern` in `text` into `out` and +/// returns how many. The pattern must not be able to overlap itself, so that every matching +/// position belongs to the non-overlapping match list; +/// [`Literal::matches_into`](crate::literal::Literal::matches_into) routes self-overlapping +/// patterns away. +/// +/// The buffer bound is asserted here, not trusted from the caller: everything unsafe in this +/// module leans on it, so the check lives next to what it protects. +pub(crate) fn matches_into( + pattern: [u8; K], + text: &[u8], + out: &mut [u32], +) -> usize { + const { + assert!( + K >= 1 && K <= 3, + "the scan is built for patterns of 1 to 3 bytes" + ); + } + assert!( + out.len() >= text.len() / K + 4, + "the match buffer must hold text.len() / pattern.len() + 4 offsets" + ); + #[cfg(target_arch = "x86_64")] + if let Some(count) = wide::matches_into_wide::(pattern, text, out) { + return count; + } + // SAFETY: the assert above is exactly `scan`'s buffer contract. + unsafe { scan::(pattern, text, out) } +} + +/// # Safety +/// `out.len() >= text.len() / K + 4` must hold; [`matches_into`] asserts it. Every match found +/// counts one slot and matches cannot overlap, so at most `text.len() / K` slots are counted +/// and the four extra keep `decode`'s unconditional writes inside the buffer. Reads from +/// `text` are justified region by region below. +unsafe fn scan(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let mut pattern_chunks = [arch::splat(0); K]; + for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) { + *chunk = arch::splat(byte); + } + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read); + // only `decode`'s debug bound check uses it. + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + + // 64 bytes per round; `any` lets a match-free round skip the masks and decodes. + // SAFETY: the loop condition keeps every read inside `text`: the furthest `match_starts` + // begins at `p + i + 48` and reads `16 + K - 1` bytes, ending at `i + 64 + (K - 1) <= len`. + // Writes are `scan`'s buffer contract, handed to `decode` as `end`. + while i + 64 + (K - 1) <= len { + let (s0, s1, s2, s3) = unsafe { + ( + arch::match_starts::(p.add(i), &pattern_chunks), + arch::match_starts::(p.add(i + 16), &pattern_chunks), + arch::match_starts::(p.add(i + 32), &pattern_chunks), + arch::match_starts::(p.add(i + 48), &pattern_chunks), + ) + }; + if arch::any(s0, s1, s2, s3) { + unsafe { + cursor = decode(arch::bits(s0), i as u32, cursor, end); + cursor = decode(arch::bits(s1), i as u32 + 16, cursor, end); + cursor = decode(arch::bits(s2), i as u32 + 32, cursor, end); + cursor = decode(arch::bits(s3), i as u32 + 48, cursor, end); + } + } + i += 64; + } + + cursor = unsafe { scan_tail::(pattern, &pattern_chunks, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } +} + +/// The remainder every scan front end shares, once its 64-byte rounds got to `i`: 16-byte +/// rounds, then one overlapped final block, or a byte-by-byte walk when the whole text is +/// shorter than a block. Returns the advanced cursor. +/// +/// # Safety +/// Same contract as [`scan`]: reads are bounded by the loop conditions, and `cursor .. end` +/// must have room for every remaining match plus `decode`'s four-slot slack. +unsafe fn scan_tail( + pattern: [u8; K], + pattern_chunks: &[arch::Chunk; K], + text: &[u8], + mut i: usize, + mut cursor: *mut u32, + end: *const u32, +) -> *mut u32 { + let len = text.len(); + let p = text.as_ptr(); + + // 16 bytes per round over what the 64-byte rounds left. + // SAFETY: reads end at `i + 16 + (K - 1) <= len`; writes as in [`scan`]. + while i + 16 + (K - 1) <= len { + cursor = unsafe { + decode( + arch::bits(arch::match_starts::(p.add(i), pattern_chunks)), + i as u32, + cursor, + end, + ) + }; + i += 16; + } + + if i + K <= len { + if len >= 16 + (K - 1) { + // One last 16-byte block, moved back to end flush with the last position a match + // can still start at. Its low bits re-cover positions the rounds above already + // decoded, so they are masked off. + // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`; + // writes as in [`scan`]. + let base = len - 16 - (K - 1); + unsafe { + let mut bits = arch::bits(arch::match_starts::(p.add(base), pattern_chunks)); + bits &= !0u64 << ((i - base) << arch::SHIFT); + cursor = decode(bits, base as u32, cursor, end); + } + } else { + // The whole text is shorter than one block: byte-by-byte costs nothing here. + while i + K <= len { + if text[i..i + K] == pattern[..] { + debug_assert!(end as usize - cursor as usize >= size_of::()); + // SAFETY: one slot per match; the buffer holds a slot for every possible + // match plus four ([`matches_into`]'s assert). + unsafe { + *cursor = i as u32; + cursor = cursor.add(1); + } + } + i += 1; + } + } + } + cursor +} diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs index 14ceed9a03..894dde46a8 100644 --- a/tokenizers/atomsplit/tests/literal.rs +++ b/tokenizers/atomsplit/tests/literal.rs @@ -48,3 +48,203 @@ fn an_empty_pattern_is_rejected() { "an empty pattern matches everywhere" ); } + +// ---- the batch scan, `matches_into` ---- + +/// Run `matches_into` with a buffer of exactly the documented size and return what it wrote. +fn batch(literal: &Literal, text: &[u8]) -> Vec { + let mut out = vec![0u32; text.len() / literal.pattern().len() + 4]; + let count = literal.matches_into(text, &mut out); + out.truncate(count); + out +} + +fn iterated(literal: &Literal, text: &[u8]) -> Vec { + literal.matches(text).map(|p| p as u32).collect() +} + +#[test] +fn matches_into_reports_the_iterator_positions() { + let dash = Literal::new(b"-").unwrap(); + assert_eq!(batch(&dash, b"a-b--c"), [1, 3, 4]); + assert_eq!(batch(&dash, b"-starts and ends-"), [0, 16]); + assert_eq!(batch(&dash, b"none here"), []); + assert_eq!(batch(&dash, b""), []); + + let space = Literal::new(b" ").unwrap(); + let text = b"the quick brown fox jumps over the lazy dog".repeat(40); + assert_eq!(batch(&space, &text), iterated(&space, &text)); + + let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap(); + let text = "\u{2581}the\u{2581}quick\u{2581}brown\u{2581}fox".repeat(40); + assert_eq!( + batch(&metaspace, text.as_bytes()), + iterated(&metaspace, text.as_bytes()) + ); +} + +/// The scan works in blocks; a match starting on either side of an internal block edge, or on +/// the last position where the pattern still fits, must not be missed or doubled. +#[test] +fn matches_into_is_seamless_across_block_boundaries() { + for pattern in [&b"-"[..], "\u{2581}".as_bytes()] { + let literal = Literal::new(pattern).unwrap(); + for start in [0, 13, 14, 15, 16, 17, 61, 62, 63, 64, 65, 127, 128] { + for len in [start + pattern.len(), 90, 129, 200] { + if start + pattern.len() > len { + continue; + } + let mut text = vec![b'x'; len]; + text[start..start + pattern.len()].copy_from_slice(pattern); + assert_eq!( + batch(&literal, &text), + [start as u32], + "pattern {pattern:?} at {start} in {len} bytes" + ); + } + } + } +} + +/// Every position matches: the largest count a text can produce, written into a buffer of +/// exactly the documented size. +#[test] +fn matches_into_handles_the_densest_text() { + let a = Literal::new(b"a").unwrap(); + let text = vec![b'a'; 300]; + assert_eq!(batch(&a, &text), (0..300).collect::>()); + + let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap(); + let text = "\u{2581}".repeat(100); + assert_eq!( + batch(&metaspace, text.as_bytes()), + (0..100).map(|i| i * 3).collect::>() + ); +} + +#[test] +fn a_self_overlapping_pattern_keeps_non_overlapping_matches() { + let aa = Literal::new(b"aa").unwrap(); + assert_eq!(batch(&aa, b"aaa"), [0]); + let run = [b'a'; 100]; + assert_eq!(batch(&aa, &run), iterated(&aa, &run)); + + let aba = Literal::new(b"aba").unwrap(); + assert_eq!(batch(&aba, b"ababa"), [0]); +} + +#[test] +fn a_pattern_longer_than_three_bytes_still_matches() { + let literal = Literal::new(b"abcd").unwrap(); + let text = b"abcd xabcdx abcabcd".repeat(20); + assert_eq!(batch(&literal, &text), iterated(&literal, &text)); +} + +#[test] +#[should_panic(expected = "matches_into")] +fn an_undersized_buffer_is_rejected() { + let a = Literal::new(b"a").unwrap(); + let mut out = vec![0u32; 7]; // "aaaa" needs 4 / 1 + 4 = 8 + a.matches_into(b"aaaa", &mut out); +} + +// ---- the streaming pair, `count_matches` + `for_each_match` ---- + +fn streamed(literal: &Literal, text: &[u8]) -> Vec { + let mut out = Vec::new(); + literal.for_each_match(text, |start| out.push(start as u32)); + out +} + +/// The streaming scan works through a fixed window; a text much longer than any plausible +/// window must report the same matches as the iterator, including around every window edge. +#[test] +fn for_each_match_streams_the_iterator_offsets_across_windows() { + for pattern in [&b" "[..], "\u{2581}".as_bytes(), b"ab"] { + let literal = Literal::new(pattern).unwrap(); + let mut text = b"word ".repeat(4_000); // ~20KB, matches every 5 bytes + text.extend_from_slice("\u{2581}ab".repeat(2_000).as_bytes()); + assert_eq!(streamed(&literal, &text), iterated(&literal, &text)); + assert_eq!( + literal.count_matches(&text), + literal.matches(&text).count(), + "count for {pattern:?}" + ); + } +} + +/// The densest possible long text: every byte starts a match, in every window. +#[test] +fn for_each_match_handles_dense_windows() { + let a = Literal::new(b"a").unwrap(); + let text = vec![b'a'; 10_000]; + assert_eq!(streamed(&a, &text), (0..10_000).collect::>()); + assert_eq!(a.count_matches(&text), 10_000); +} + +#[test] +fn for_each_match_takes_the_iterator_route_for_uncovered_patterns() { + // Self-overlapping and longer-than-three patterns are exactly the ones the batch scan + // refuses; the streaming pair must still report the iterator's non-overlapping matches. + for pattern in [&b"aa"[..], b"aba", b"abcd"] { + let literal = Literal::new(pattern).unwrap(); + let text = b"aabaabcdaaaba".repeat(2_000); + assert_eq!(streamed(&literal, &text), iterated(&literal, &text)); + assert_eq!(literal.count_matches(&text), literal.matches(&text).count()); + } +} + +#[test] +fn streaming_agrees_with_the_iterator_on_random_lengths() { + let mut state = 0x1234_5678_9ABC_DEF0u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let patterns: &[&[u8]] = &[b"a", b"ab", b"aa", b"aba", b"abc", b"abca"]; + for _ in 0..500 { + // lengths spread around plausible window sizes, so edges get hit both exactly and off by one + let len = (next() % 8_000) as usize; + let text: Vec = (0..len).map(|_| b'a' + (next() % 3) as u8).collect(); + let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap(); + assert_eq!( + streamed(&literal, &text), + iterated(&literal, &text), + "pattern {:?} len {len}", + literal.pattern() + ); + assert_eq!(literal.count_matches(&text), literal.matches(&text).count()); + } +} + +#[test] +fn matches_into_agrees_with_the_iterator_on_random_inputs() { + // xorshift64: deterministic, no dev-dependency needed. + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let patterns: &[&[u8]] = &[ + b"a", b"b", b"ab", b"aa", b"ba", b"aab", b"aba", b"abc", b"abca", + ]; + for _ in 0..20_000 { + let len = (next() % 300) as usize; + let alphabet = 2 + (next() % 2) as u8; + let text: Vec = (0..len) + .map(|_| b'a' + (next() % alphabet as u64) as u8) + .collect(); + let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap(); + assert_eq!( + batch(&literal, &text), + iterated(&literal, &text), + "pattern {:?} text {:?}", + literal.pattern(), + text + ); + } +} diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index cc4545b47a..e23d20399d 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -145,10 +145,22 @@ impl pipeline::Normalizer for Replace { Ok(match &self.search { Search::Literal(literal) => { let width = literal.pattern().len(); - let matches = literal - .matches(input.as_bytes()) - .map(|start| (start, start + width)); - replace_matches(input, &self.content, matches) + let count = literal.count_matches(input.as_bytes()); + if count == 0 { + return Ok(Cow::Borrowed(input)); + } + // Counting first buys the exact output length (every match swaps `width` + // bytes for the content), so the build below never reallocates. + let exact = input.len() - count * width + count * self.content.len(); + let mut replaced = String::with_capacity(exact); + let mut last_end = 0; + literal.for_each_match(input.as_bytes(), |start| { + replaced.push_str(&input[last_end..start]); + replaced.push_str(&self.content); + last_end = start + width; + }); + replaced.push_str(&input[last_end..]); + Cow::Owned(replaced) } Search::Regex(regex) => replace_matches(input, &self.content, regex.find_iter(input)), Search::Nothing => Cow::Borrowed(input), diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 779c167f10..bc5682cadd 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -201,13 +201,33 @@ impl pipeline::PreTokenizer for Split { ); return Ok(()); } - // Not a natively-routed GPT regex: fall-back to Literal or Regex search + // A plain-string pattern streams: the batch scan hands each delimiter offset straight + // to the span fold, and nothing is built in between. The count pass sizes `out` for + // the worst case, one span per match plus the pieces around them. + if let Search::Literal(literal) = &self.search { + let width = literal.pattern().len(); + let count = literal.count_matches(text.as_bytes()); + out.reserve(2 * count + 1); + let mut fold = pipeline::SplitFold::new(out, self.behavior); + let mut prev = 0; + literal.for_each_match(text.as_bytes(), |start| { + if prev != start { + fold.segment((prev, start), self.invert); + } + fold.segment((start, start + width), !self.invert); + prev = start + width; + }); + if prev != text.len() { + fold.segment((prev, text.len()), self.invert); + } + fold.finish(); + return Ok(()); + } + // Not a natively-routed GPT regex either: fall back to the Regex search let matches = match (&self.search, self.invert) { - (Search::Literal(literal), false) => literal.find_matches(text)?, - (Search::Literal(literal), true) => Invert(literal).find_matches(text)?, (Search::Regex(regex), false) => regex.find_matches(text)?, (Search::Regex(regex), true) => Invert(regex).find_matches(text)?, - (Search::Unavailable, _) => { + _ => { return Err( "this `Split` pattern needs a system-regex backend; enable the `fancy-regex` feature" .into(), @@ -415,6 +435,84 @@ mod tests { ); } + /// The literal path across every behavior, on a text with a delimiter at both ends and a + /// consecutive pair — the shapes the span fold can get wrong. + #[test] + fn pipeline_literal_covers_every_behavior() { + let text = "-a--b-"; + #[allow(clippy::type_complexity)] + let cases: Vec<(SplitDelimiterBehavior, Vec<(&str, (u32, u32))>)> = vec![ + (Removed, vec![("a", (1, 2)), ("b", (4, 5))]), + ( + Isolated, + vec![ + ("-", (0, 1)), + ("a", (1, 2)), + ("-", (2, 3)), + ("-", (3, 4)), + ("b", (4, 5)), + ("-", (5, 6)), + ], + ), + ( + MergedWithPrevious, + vec![("-", (0, 1)), ("a-", (1, 3)), ("-", (3, 4)), ("b-", (4, 6))], + ), + ( + MergedWithNext, + vec![("-a", (0, 2)), ("-", (2, 3)), ("-b", (3, 5)), ("-", (5, 6))], + ), + ( + Contiguous, + vec![ + ("-", (0, 1)), + ("a", (1, 2)), + ("--", (2, 4)), + ("b", (4, 5)), + ("-", (5, 6)), + ], + ), + ]; + for (behavior, expected) in cases { + assert_eq!( + pipeline_split("-".into(), behavior, false, text), + expected, + "behavior: {behavior:?}", + ); + } + } + + /// Inverted literal search: the gaps between delimiters become the matches. + #[test] + fn pipeline_literal_inverts() { + // Removed drops the (inverted) matches, keeping each delimiter. + assert_eq!( + pipeline_split("-".into(), Removed, true, "-a--b-"), + vec![("-", (0, 1)), ("-", (2, 3)), ("-", (3, 4)), ("-", (5, 6))], + ); + // Isolated keeps everything either way. + assert_eq!( + pipeline_split("-".into(), Isolated, true, "a-b"), + vec![("a", (0, 1)), ("-", (1, 2)), ("b", (2, 3))], + ); + } + + #[test] + fn pipeline_literal_edges() { + assert_eq!( + pipeline_split("-".into(), Removed, false, ""), + Vec::<(&str, (u32, u32))>::new(), + ); + assert_eq!( + pipeline_split("-".into(), Removed, false, "---"), + Vec::<(&str, (u32, u32))>::new(), + ); + assert_eq!( + pipeline_split("-".into(), Removed, false, "abc"), + vec![("abc", (0, 3))], + ); + } + #[test] fn pipeline_gpt2_uses_fsm_and_matches_legacy() { // The gpt2 pattern is recognized -> the pipeline path routes to the native diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5a147f5f4a..a411ef2657 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -43,22 +43,27 @@ impl Pattern for &Regex { } /// Searching for a plain string, does not need a regex engine: [`Literal`] scans the bytes. +/// Two batch-scan passes: [`Literal::count_matches`] sizes `splits` exactly (each match adds +/// itself and at most one gap before it, plus the final gap), then +/// [`Literal::for_each_match`] streams the offsets in. impl Pattern for &Literal { fn find_matches(&self, inside: &str) -> Result> { if inside.is_empty() { return Ok(vec![((0, 0), false)]); } + let width = self.pattern().len(); + let count = self.count_matches(inside.as_bytes()); + let mut splits = Vec::with_capacity(2 * count + 1); let mut prev = 0; - let mut splits = Vec::with_capacity(inside.len()); - for start in self.matches(inside.as_bytes()) { - let end = start + self.pattern().len(); + self.for_each_match(inside.as_bytes(), |start| { + let end = start + width; if prev != start { splits.push(((prev, start), false)); } splits.push(((start, end), true)); prev = end; - } + }); if prev != inside.len() { splits.push(((prev, inside.len()), false)) } diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 01e2172961..d55cc7a139 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -982,75 +982,101 @@ pub fn split_matches( matches: Vec<((usize, usize), bool)>, behavior: SplitDelimiterBehavior, ) { - use SplitDelimiterBehavior::*; + let mut fold = SplitFold::new(out, behavior); + for (offsets, is_match) in matches { + fold.segment(offsets, is_match); + } + fold.finish(); +} - // (offsets, should_remove), mirroring `NormalizedString::split`. - let splits: Vec<((usize, usize), bool)> = match behavior { - Isolated => matches.into_iter().map(|(o, _)| (o, false)).collect(), - Removed => matches, // should_remove == is_match - Contiguous => { - let mut previous_match = false; - matches - .into_iter() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match == previous_match { - if let Some(((_, end), _)) = acc.last_mut() { - *end = offsets.1; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }) +/// The streaming form of [`split_matches`]: feed it the covering `(offsets, is_match)` +/// segments left to right with [`SplitFold::segment`], then call [`SplitFold::finish`]. +/// Callers that produce segments one at a time (the literal path in `Split`) drive it +/// directly and never build the segment vector. +/// +/// One span under construction and the previous segment's flag are the only state. A span +/// stays pending as long as the next segment might still extend it (a delimiter merging with +/// its neighbour, a contiguous run); a segment that cannot extend it flushes it to `out` and +/// takes its place. +pub(crate) struct SplitFold<'a> { + out: &'a mut Vec, + behavior: SplitDelimiterBehavior, + pending: Option<(usize, usize)>, + previous_match: bool, +} + +impl<'a> SplitFold<'a> { + pub(crate) fn new(out: &'a mut Vec, behavior: SplitDelimiterBehavior) -> Self { + Self { + out, + behavior, + pending: None, + previous_match: false, } - MergedWithPrevious => { - let mut previous_match = false; - matches - .into_iter() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match && !previous_match { - if let Some(((_, end), _)) = acc.last_mut() { - *end = offsets.1; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }) + } + + pub(crate) fn segment(&mut self, offsets: (usize, usize), is_match: bool) { + use SplitDelimiterBehavior::*; + match self.behavior { + Isolated => self.emit(Some(offsets)), + Removed => { + if !is_match { + self.emit(Some(offsets)); + } + } + // A run of equal flags becomes one span. + Contiguous => { + if is_match == self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } + // A delimiter glues onto the piece before it; in a run of delimiters only the + // first one glues, the rest stand alone. + MergedWithPrevious => { + if is_match && !self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } + // A delimiter glues onto the piece after it, so it stays pending until that + // piece arrives; a delimiter followed by another delimiter stands alone. + MergedWithNext => { + if !is_match && self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } } - MergedWithNext => { - let mut previous_match = false; - let mut splits = - matches - .into_iter() - .rev() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match && !previous_match { - if let Some(((start, _), _)) = acc.last_mut() { - *start = offsets.0; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }); - splits.reverse(); - splits + self.previous_match = is_match; + } + + /// Flushes the last pending span. + pub(crate) fn finish(mut self) { + let last = self.pending.take(); + self.emit(last); + } + + /// Stretches the pending span to cover `offsets` too, or starts one from it. + fn extend_pending_to(&mut self, offsets: (usize, usize)) { + match &mut self.pending { + Some((_, end)) => *end = offsets.1, + None => self.pending = Some(offsets), } - }; + } - for ((start, end), should_remove) in splits { - if !should_remove && start != end { - out.push(Span { + /// Empty pieces are dropped, matching the fold in `NormalizedString::split`. + fn emit(&mut self, span: Option<(usize, usize)>) { + if let Some((start, end)) = span + && start != end + { + self.out.push(Span { start: start as u32, end: end as u32, }); From 9769d6790242d965a76f4369539337de9a07532e Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:37:40 +0200 Subject: [PATCH 03/13] perf(bpe): cut whole-chunk SentencePiece text into words, proven from the vocabulary (#2266) SentencePiece models hand the model one whole chunk. Where the vocabulary proves no merge can cross a word boundary, cut there and tokenize the words independently. Squashed from perf/metaspace-proven-cuts. --- tokenizers/tk-encode/src/models/bpe/model.rs | 11 + .../tk-encode/src/normalizers/replace.rs | 5 + .../tk-encode/src/pre_tokenizers/mod.rs | 1 + .../src/pre_tokenizers/proven_cuts.rs | 641 ++++++++++++++++++ .../tk-encode/src/tokenizer/pipeline.rs | 17 + tokenizers/tk-encode/tests/pipeline_oracle.rs | 74 +- 6 files changed, 723 insertions(+), 26 deletions(-) create mode 100644 tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 9e4853e9b7..dbcfb254c1 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -767,6 +767,17 @@ impl PipelineBPE { }) } + /// Does this model look the whole text up in the vocabulary before merging it? + pub(crate) fn ignore_merges(&self) -> bool { + self.ignore_merges + } + + /// Every piece of the vocabulary, as bytes. Used at build time to work out how the text may be + /// cut into words; not cheap, so call it once. + pub(crate) fn vocab_bytes(&self) -> Vec<(Vec, u32)> { + self.vocab.byte_content() + } + fn merge_word( &self, sequence: &str, diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index e23d20399d..03e207360d 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -103,6 +103,11 @@ impl Replace { search, }) } + + /// What this normalizer looks for. + pub fn pattern(&self) -> &ReplacePattern { + &self.pattern + } } impl Normalizer for Replace { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/mod.rs b/tokenizers/tk-encode/src/pre_tokenizers/mod.rs index b895e3f294..757cf4ceb9 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/mod.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/mod.rs @@ -4,6 +4,7 @@ pub mod delimiter; pub mod digits; pub mod fixed_length; pub mod metaspace; +pub mod proven_cuts; pub mod punctuation; pub mod sequence; pub mod split; diff --git a/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs new file mode 100644 index 0000000000..2ad2ee4047 --- /dev/null +++ b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs @@ -0,0 +1,641 @@ +//! Cutting a whole SentencePiece text into words, where the vocabulary proves it is harmless. +//! +//! SentencePiece vocabularies (gemma, llama-2, …) write a space as `▁` (U+2581), so `"tell me"` +//! reaches the model as `"▁tell▁me"`. That character both opens a word and separates two, which is +//! why the code calls it the delimiter. +//! +//! Some of these tokenizers ship a [`Metaspace`] pre-tokenizer, which writes the delimiters and cuts +//! before every one of them. Those cuts are the tokenizer's own output, so +//! [`super::metaspace::to_normalizer_and_split`] reproduces them exactly and there is nothing to +//! decide. +//! +//! The tokenizers here are the other kind: their normalizer writes the delimiters and they ship no +//! pre-tokenizer that cuts, so the model receives the whole text in one piece. Cutting it into words +//! is only a speed-up — merging a long text costs more than merging its words one after another, and +//! short words are friendlier to the cache. But a speed-up is worthless if it changes the tokens, so +//! here every cut has to be proven harmless first. Two rules do that: +//! +//! 1. A group of delimiters stays with the word that follows it (`a▁▁▁b` → `a`, `▁▁▁b`). +//! Vocabularies hold pieces made of several delimiters (`▁▁`, `▁▁▁`), and cutting inside a group +//! would stop those from forming. +//! 2. A cut is dropped when a vocabulary piece could merge across it. A merge only ever produces a +//! piece that is in the vocabulary, so this can only happen if some piece holds a delimiter that +//! is not at its start. gemma has exactly one such piece: `>▁ bool { + self.delimiter.pattern() == other.delimiter.pattern() && self.veto == other.veto + } +} + +impl ProvenCuts { + fn new(delimiter: Literal, veto: Veto) -> Self { + Self { delimiter, veto } + } +} + +impl pipeline::PreTokenizer for ProvenCuts { + fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { + let delimiter = self.delimiter.pattern(); + let bytes = text.as_bytes(); + // The word being built runs from `start` to the next delimiter we cut at. + let mut start = 0usize; + for at in self.delimiter.matches(bytes) { + // A delimiter at `start` is the cut we just made, or the text opens with one: either way + // the word would be empty. + if at == start { + continue; + } + // A delimiter right in front of this one means we are inside a group of them, and a group + // stays whole with the word after it. + if bytes[..at].ends_with(delimiter) || self.veto.forbids(bytes, at) { + continue; + } + out.push(Span::new(start as u32, at as u32)); + start = at; + } + if start < bytes.len() { + out.push(Span::new(start as u32, bytes.len() as u32)); + } + Ok(()) + } +} + +/// The splitter for a tokenizer whose text reaches the model in one piece, or `None` when the text +/// has to stay that way — because something already cuts it, because no normalizer provably writes +/// the delimiters, or because the vocabulary cannot prove the cuts. +pub(crate) fn for_tokenizer( + normalizer: Option<&NormalizerWrapper>, + pre_tokenizer: Option<&PreTokenizerWrapper>, + model: &PipelineModel, +) -> Option { + if !leaves_the_text_whole(pre_tokenizer) { + return None; + } + let delimiter = delimiter_from_normalizer(normalizer?)?; + let delimiter = Literal::new(delimiter.to_string().as_bytes()).expect("a char is never empty"); + let veto = veto_from_model(model, &delimiter)?; + Some(ProvenCuts::new(delimiter, veto)) +} + +/// Does this pre-tokenizer leave the text in one piece? Only the two shapes below do, and anything +/// else is refused rather than guessed at: cutting text a tokenizer meant to keep whole changes the +/// tokens it produces. +fn leaves_the_text_whole(pre_tokenizer: Option<&PreTokenizerWrapper>) -> bool { + match pre_tokenizer { + // Nothing cuts the text at all. llama-2 ships this shape. + None => true, + // A `Split` that can never match, because the normalizer already replaced every space it + // looks for. With nothing to match, every behaviour it could carry leaves the text in one + // piece. gemma ships this shape. + Some(PreTokenizerWrapper::Split(split)) => { + !split.invert + && matches!(&split.pattern, SplitPattern::String(pattern) + if !pattern.is_empty() && pattern.chars().all(|c| c == ' ')) + } + _ => false, + } +} + +/// The cuts this model's vocabulary forbids, or `None` when it cannot be cut at all. +fn veto_from_model(model: &PipelineModel, delimiter: &Literal) -> Option { + let PipelineModel::BPE(bpe) = model else { + return None; + }; + // With `ignore_merges` the model first looks the whole text up in the vocabulary and emits one + // token when it finds it. Handing it words instead would skip that lookup. + if bpe.ignore_merges() { + return None; + } + Veto::build(&bpe.vocab_bytes(), delimiter) +} + +/// The character every space becomes, if the normalizer provably rewrites all of them. +/// +/// Accepts a `Replace` on its own, or a sequence whose last step is one and whose earlier steps only +/// prepend text — a step running after the `Replace` could bring spaces back. +fn delimiter_from_normalizer(normalizer: &NormalizerWrapper) -> Option { + match normalizer { + NormalizerWrapper::Replace(replace) => space_replacement(replace), + NormalizerWrapper::Sequence(sequence) => { + let (last, rest) = sequence.as_ref().split_last()?; + rest.iter() + .all(|step| matches!(step, NormalizerWrapper::Prepend(_))) + .then(|| match last { + NormalizerWrapper::Replace(replace) => space_replacement(replace), + _ => None, + }) + .flatten() + } + _ => None, + } +} + +/// The single character this normalizer turns every space into. +fn space_replacement(replace: &Replace) -> Option { + if replace.pattern() != &ReplacePattern::String(" ".to_string()) { + return None; + } + let mut content = replace.content.chars(); + let delimiter = content.next()?; + content.next().is_none().then_some(delimiter) +} + +/// How many veto pieces we put up with before giving up on cutting at all. +const MAX_VETO_PIECES: usize = 32; + +/// Width of one padded half of a veto piece. Whatever length the halves really are, they are compared +/// a full `u128` at a time. +const HALF_WIDTH: usize = size_of::(); + +/// One half of a [`VetoPiece`], padded out to [`HALF_WIDTH`] bytes. +/// +/// The halves are a byte or two long, but their length is only known once the vocabulary is read, and +/// comparing a run-time number of bytes means calling `memcmp`. Padding to a fixed width instead +/// turns the compare into a couple of register operations, which is worth it in a loop that runs once +/// per word. +#[derive(Debug, Clone, PartialEq, Eq)] +struct Half { + bytes: u128, + /// `0xff` over the bytes of the half, `0` over the padding. + mask: u128, +} + +impl Half { + /// The half at the end of the window, where the bytes running up to a cut land. + fn ending(half: &[u8]) -> Option { + Self::padded(half, HALF_WIDTH.checked_sub(half.len())?) + } + + /// The half at the start of the window, where the bytes following a cut land. + fn starting(half: &[u8]) -> Option { + Self::padded(half, 0) + } + + /// `None` when the half is wider than the window. + fn padded(half: &[u8], at: usize) -> Option { + let mut bytes = [0u8; HALF_WIDTH]; + let mut mask = [0u8; HALF_WIDTH]; + bytes.get_mut(at..at + half.len())?.copy_from_slice(half); + mask[at..at + half.len()].fill(0xff); + Some(Self { + bytes: u128::from_le_bytes(bytes), + mask: u128::from_le_bytes(mask), + }) + } + + fn matches(&self, window: u128) -> bool { + (window ^ self.bytes) & self.mask == 0 + } +} + +/// A vocabulary piece holding a delimiter that is not at its start, split at that delimiter: `>▁", after: " Option { + Some(Self { + before: Half::ending(before)?, + after: Half::starting(after)?, + }) + } +} + +/// The cuts a vocabulary does not allow. +/// +/// A merge only ever produces a piece that is in the vocabulary, so a merge can only reach across a +/// cut if some piece holds a delimiter that is not at its start. Those pieces are collected here, and +/// a cut where one of them fits is dropped. +#[derive(Debug, Clone, PartialEq, Eq)] +struct Veto { + /// Empty means no merge can reach across a cut, so every cut is allowed. + pieces: Vec, + /// Last bytes of the `before` halves, as a 256-bit set. A cut with any other byte in front of it + /// fits no piece, which is how most cuts get away without a single comparison. + bytes_before: [u64; 4], + /// Length of the delimiter the pieces were split at. Every `after` half starts at the byte right + /// behind the delimiter, so the window compared to it has to start there too. + delimiter_len: usize, +} + +impl Veto { + /// Reads the vocabulary and collects the pieces a merge could use to reach across a cut. + /// + /// `None` turns cutting off: either this is not a SentencePiece vocabulary (no delimiter piece in + /// it), or its veto pieces are too many, too long or too tangled to rule out cheaply. + fn build(vocab: &[(Vec, u32)], delimiter: &Literal) -> Option { + let pattern = delimiter.pattern(); + if !vocab.iter().any(|(piece, _)| piece.as_slice() == pattern) { + return None; + } + let mut pieces = Vec::new(); + let mut bytes_before = [0u64; 4]; + for (piece, _) in vocab { + for at in delimiter.matches(piece) { + // A piece starting with a delimiter sits right after a cut, not across it, and a + // delimiter following another one is inside a group, where we never cut. + if at == 0 || piece[..at].ends_with(pattern) { + continue; + } + let (before, after) = (&piece[..at], &piece[at + pattern.len()..]); + // A piece with a second delimiter reaches across two cuts at once, and checking one + // cut at a time no longer proves anything. None of the vocabularies we tested has + // one, so drop cutting instead. + if pieces.len() == MAX_VETO_PIECES + || delimiter.matches(before).next().is_some() + || delimiter.matches(after).next().is_some() + { + return None; + } + let previous = *before.last().expect("`at` is past the start of the piece"); + bytes_before[(previous >> 6) as usize] |= 1 << (previous & 63); + pieces.push(VetoPiece::new(before, after)?); + } + } + Some(Self { + pieces, + bytes_before, + delimiter_len: pattern.len(), + }) + } + + /// Could a piece cover the delimiter at `at`, so that a merge reaches over a cut placed there? A + /// match only means such a merge is possible, not that the model performs it — either way we + /// leave the text in one piece. `at` is past the start of the text, so there is a byte in front + /// of it. + fn forbids(&self, text: &[u8], at: usize) -> bool { + let previous = text[at - 1]; + if self.bytes_before[(previous >> 6) as usize] & (1 << (previous & 63)) == 0 { + return false; + } + // Both windows are padded with zeros, so a half can only match beyond the ends of `text` if + // the half itself holds a zero byte — and one match too many only leaves the text uncut. + let before = window_ending(&text[..at]); + let after = window_starting(&text[at + self.delimiter_len..]); + self.pieces + .iter() + .any(|piece| piece.before.matches(before) && piece.after.matches(after)) + } +} + +/// The last [`HALF_WIDTH`] bytes of `bytes`, padded on the left when there are fewer, packed the way +/// [`Half::ending`] packs a half. +fn window_ending(bytes: &[u8]) -> u128 { + match bytes.last_chunk() { + Some(window) => u128::from_le_bytes(*window), + None => { + let mut window = [0u8; HALF_WIDTH]; + window[HALF_WIDTH - bytes.len()..].copy_from_slice(bytes); + u128::from_le_bytes(window) + } + } +} + +/// The first [`HALF_WIDTH`] bytes of `bytes`, padded on the right when there are fewer, packed the way +/// [`Half::starting`] packs a half. +fn window_starting(bytes: &[u8]) -> u128 { + match bytes.first_chunk() { + Some(window) => u128::from_le_bytes(*window), + None => { + let mut window = [0u8; HALF_WIDTH]; + window[..bytes.len()].copy_from_slice(bytes); + u128::from_le_bytes(window) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::models::bpe::{BPE, BpeBuilder, Merges, PipelineBPE, Vocab}; + use crate::tokenizer::Model; + + const DELIMITER: char = '▁'; + + fn delimiter() -> Literal { + Literal::new(DELIMITER.to_string().as_bytes()).unwrap() + } + + fn split_on(veto: Veto, text: &str) -> Vec { + let split = ProvenCuts::new(delimiter(), veto); + let mut spans = Vec::new(); + pipeline::PreTokenizer::pre_tokenize(&split, text, &mut spans).unwrap(); + spans.iter().map(|s| text[s.range()].to_string()).collect() + } + + fn veto_of(vocab: &[&str]) -> Option { + let pieces: Vec<(Vec, u32)> = vocab + .iter() + .enumerate() + .map(|(id, piece)| (piece.as_bytes().to_vec(), id as u32)) + .collect(); + Veto::build(&pieces, &delimiter()) + } + + /// A vocabulary shaped like gemma's: `>▁ BPE { + let vocab: Vocab = [ + ("▁", 0u32), + ("<", 1), + (">", 2), + ("/", 3), + ("s", 4), + ("p", 5), + ("a", 6), + ("b", 7), + ("▁", 9), + (">▁", "▁"), + (">▁", " Vec { + let mut spans = Vec::new(); + pipeline::PreTokenizer::pre_tokenize(split, text, &mut spans).unwrap(); + let mut out = Vec::new(); + let mut scratch = pipeline::Model::init_scratch(model); + for span in &spans { + pipeline::Model::tokenize_pipeline(model, &text[span.range()], &mut scratch, &mut out) + .unwrap(); + } + out.iter().map(|token| token.id).collect() + } + + /// The ids the model produces for the whole text at once — what the cuts must reproduce. + fn ids_whole_text(model: &BPE, text: &str) -> Vec { + model + .tokenize(text) + .unwrap() + .iter() + .map(|token| token.id) + .collect() + } + + #[test] + fn proven_cuts_keep_the_whole_text_ids() { + let reference = metaspace_bpe(); + let model = PipelineBPE::from_bpe(reference.clone(), false).unwrap(); + let veto = Veto::build(&model.vocab_bytes(), &delimiter()).unwrap(); + let split = ProvenCuts::new(delimiter(), veto); + for text in [ + "▁a▁b", + "a▁b", + "▁▁▁a", + "a▁▁▁b", + "▁sp▁a", + "sp", + "a▁", + "▁", + "▁a▁sp▁b▁a", + // The veto pieces, alone and surrounded by other cuts. + "▁", + "▁a▁▁b", + "p▁a", + "▁ap▁a▁b", + // Starts like a veto piece but ends differently, so the cut stands. + "▁a", + "q▁a", + ] { + assert_eq!( + ids_word_by_word(&model, &split, text), + ids_whole_text(&reference, text), + "{text:?}" + ); + } + } + + #[test] + fn proven_cuts_keep_delimiter_groups_whole() { + let veto = || veto_of(&["▁"]).unwrap(); + assert_eq!(split_on(veto(), "a▁▁▁b"), ["a", "▁▁▁b"]); + assert_eq!(split_on(veto(), "▁▁▁a"), ["▁▁▁a"]); + assert_eq!(split_on(veto(), "hello"), ["hello"]); + assert_eq!(split_on(veto(), ""), Vec::::new()); + } + + #[test] + fn proven_cuts_skip_what_the_veto_forbids() { + let veto = || veto_of(&["▁", ">▁▁▁"), ["▁"]); + // ">" in front of the `▁` again, but "▁a"), ["", "▁a"]); + // Only the cut the piece covers is dropped, not the other ones. + assert_eq!(split_on(veto(), "x▁▁▁y"), ["x", "▁▁", "▁y"]); + } + + #[test] + fn veto_needs_the_delimiter_in_the_vocabulary() { + assert!(veto_of(&["a", "b"]).is_none()); + assert!(veto_of(&["a", "▁"]).is_some()); + } + + #[test] + fn veto_collects_only_the_pieces_that_reach_over_a_cut() { + assert!(veto_of(&["▁", "▁hello", "▁▁"]).unwrap().pieces.is_empty()); + assert_eq!( + veto_of(&["▁", "p▁"]).unwrap().pieces, + [VetoPiece::new(b"p", b"").unwrap()] + ); + } + + #[test] + fn veto_gives_up_when_it_cannot_prove_anything() { + // Two delimiters in one piece would reach over two cuts at once. + assert!(veto_of(&["▁", "a▁b▁c"]).is_none()); + // More pieces than the loop is willing to walk per cut. + let mut vocab = vec!["▁".to_string()]; + vocab.extend((0..=MAX_VETO_PIECES).map(|i| format!("{i}▁x"))); + let vocab: Vec<&str> = vocab.iter().map(String::as_str).collect(); + assert!(veto_of(&vocab).is_none()); + // A half that does not fit the fixed-width compare. + let wide = "a".repeat(HALF_WIDTH + 1); + assert!(veto_of(&["▁", &format!("{wide}▁x")]).is_none()); + assert!(veto_of(&["▁", &format!("a▁{wide}")]).is_none()); + } + + #[test] + fn veto_matches_a_half_that_fills_the_window() { + let before = "a".repeat(HALF_WIDTH); + let veto = || veto_of(&["▁", &format!("{before}▁x"), "▁x"]).unwrap(); + let text = format!("{before}▁x"); + assert_eq!(split_on(veto(), &text), [text.as_str()]); + // One byte short of the half: the window pads with a zero the half does not hold, so the + // piece cannot form and the cut stands. + let short = &before[1..]; + assert_eq!(split_on(veto(), &format!("{short}▁x")), [short, "▁x"]); + } + + mod for_tokenizer { + use super::*; + + /// The shapes a `▁`-spelling tokenizer ships, copied from the files in `data/`. + const GEMMA_NORMALIZER: &str = + r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#; + const GEMMA_PRE_TOKENIZER: &str = r#"{"type":"Split","pattern":{"String":" "},"behavior":"MergedWithPrevious","invert":false}"#; + const LLAMA_NORMALIZER: &str = r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#; + + fn normalizer(json: &str) -> NormalizerWrapper { + serde_json::from_str(json).unwrap() + } + + fn pre_tokenizer(json: &str) -> PreTokenizerWrapper { + serde_json::from_str(json).unwrap() + } + + fn bpe_model(bpe: BPE) -> PipelineModel { + PipelineModel::BPE(PipelineBPE::from_bpe(bpe, false).unwrap()) + } + + #[test] + fn the_shapes_that_leave_the_text_whole_are_cut() { + let model = bpe_model(metaspace_bpe()); + for (name, normalizer_json, pre_tokenizer_json) in [ + ("gemma", GEMMA_NORMALIZER, Some(GEMMA_PRE_TOKENIZER)), + ("llama-2", LLAMA_NORMALIZER, None), + ] { + let cuts = for_tokenizer( + Some(&normalizer(normalizer_json)), + pre_tokenizer_json.map(pre_tokenizer).as_ref(), + &model, + ) + .unwrap_or_else(|| panic!("{name} should be cut into words")); + assert_eq!(cuts.delimiter.pattern(), "▁".as_bytes(), "{name}"); + } + } + + #[test] + fn refuses_what_it_cannot_prove() { + let model = bpe_model(metaspace_bpe()); + let refused: [(&str, Option<&NormalizerWrapper>, Option<&str>); 5] = [ + ("no normalizer and no pre-tokenizer", None, None), + // Nothing here turns spaces into the delimiter. + ( + "a normalizer that keeps spaces", + Some(&normalizer(r#"{"type":"NFC"}"#)), + None, + ), + // A step running after the replace could bring spaces back. + ( + "a replace that is not the last step", + Some(&normalizer( + r#"{"type":"Sequence","normalizers":[{"type":"Replace","pattern":{"String":" "},"content":"▁"},{"type":"NFC"}]}"#, + )), + None, + ), + // This one cuts the text itself, so the model never sees it whole. + ( + "a split that does match", + Some(&normalizer(GEMMA_NORMALIZER)), + Some( + r#"{"type":"Split","pattern":{"String":"▁"},"behavior":"MergedWithNext","invert":false}"#, + ), + ), + ( + "a byte-level pre-tokenizer", + Some(&normalizer(GEMMA_NORMALIZER)), + Some( + r#"{"type":"ByteLevel","add_prefix_space":false,"trim_offsets":true,"use_regex":true}"#, + ), + ), + ]; + for (name, normalizer, json) in refused { + let declared = json.map(pre_tokenizer); + assert!( + for_tokenizer(normalizer, declared.as_ref(), &model).is_none(), + "{name}" + ); + } + // A vocabulary without the delimiter proves nothing about cutting on it. + let plain: Vocab = [("a", 0u32), ("b", 1)] + .iter() + .map(|(piece, id)| ((*piece).into(), *id)) + .collect(); + let plain = BpeBuilder::default() + .vocab_and_merges(plain, Vec::new()) + .build() + .unwrap(); + assert!( + for_tokenizer(Some(&normalizer(GEMMA_NORMALIZER)), None, &bpe_model(plain)) + .is_none(), + "vocabulary without the delimiter" + ); + } + + /// The real files, so a change to either config shape shows up here. The veto counts are the + /// premise the whole proof rests on — gemma-4 has one piece that can merge across a word + /// boundary (`>▁ pretok.pre_tokenize(text, out), Self::Digits(pretok) => pretok.pre_tokenize(text, out), Self::FixedLength(pretok) => pretok.pre_tokenize(text, out), + Self::ProvenCuts(pretok) => pretok.pre_tokenize(text, out), Self::Punctuation(pretok) => pretok.pre_tokenize(text, out), Self::Sequence(pretok) => pretok.pre_tokenize(text, out), Self::Split(pretok) => pretok.pre_tokenize(text, out), @@ -633,6 +636,16 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { ModelWrapper::WordPiece(model) => PipelineModel::WordPiece(model.try_into()?), }; + // Some tokenizers spell every space as `▁` and then cut nowhere, so the model receives the + // whole text in one piece. Cutting it into words is faster wherever the vocabulary proves + // that leaves the ids alone — and that proof reads the vocabulary, hence after the model. + let pre_tokenizer = + match proven_cuts::for_tokenizer(tok.get_normalizer(), tok.get_pre_tokenizer(), &model) + { + Some(cuts) => PipelinePreTokenizer::ProvenCuts(cuts), + None => pre_tokenizer, + }; + Ok(Self { added_vocabulary, normalizers, @@ -663,6 +676,10 @@ impl PipelineTokenizer { &self.model } + pub fn get_pre_tokenizer(&self) -> &PipelinePreTokenizer { + &self.pre_tokenizer + } + /// Encode `input` into token ids. /// /// Special tokens are matched in two passes: diff --git a/tokenizers/tk-encode/tests/pipeline_oracle.rs b/tokenizers/tk-encode/tests/pipeline_oracle.rs index 6222084627..1582c2f98c 100644 --- a/tokenizers/tk-encode/tests/pipeline_oracle.rs +++ b/tokenizers/tk-encode/tests/pipeline_oracle.rs @@ -26,6 +26,20 @@ use tokenizers_release::Tokenizer as Released; const PROBE: &str = "The quick brown fox jumps 123."; +/// Inputs built to break word boundaries, which the natural corpora do not reach. ` ` is the +/// one that matters most: gemma-4 merges its `>`, mark and ` crossing > one

two

", + "The quick brown fox \u{2014} jumps; over, the: lazy. dog! ", + "\t\tindent()\twide spacing here \r\nCRLF\r\n", + "\u{2581}\u{2581}\u{2581}already marked\u{2581}", + "trailing space and mark p\u{2581}q ", + "\u{65e5}\u{672c}\u{8a9e}\u{306e}\u{30c6}\u{30ad}\u{30b9}\u{30c8} \u{1f389} combining a\u{301}e\u{308} nbsp\u{a0}here", + " leading and trailing ", +]; + fn check_model(tok_file: &str) { let path = Path::new(DATA).join(tok_file); // The legacy `Tokenizer` only *builds* the pipeline (its sole constructor @@ -52,6 +66,39 @@ fn check_model(tok_file: &str) { }; let mut failures = Vec::new(); + let mut check = |case: String, chunk: &str| { + for add_special_tokens in [false, true] { + let expected = released + .encode_fast(chunk, add_special_tokens) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipeline + .encode(chunk, add_special_tokens) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + if expected != got { + let at = expected + .iter() + .zip(&got) + .position(|(e, g)| e != g) + .unwrap_or(expected.len().min(got.len())); + failures.push(format!( + "{case} (add_special_tokens={add_special_tokens}): ids diverge at {at} \ + (expected len {}, got len {})", + expected.len(), + got.len(), + )); + } + } + }; + + for (i, chunk) in HOSTILE.iter().enumerate() { + check(format!("hostile/{i}"), chunk); + } + for &(group, stem) in FIXTURES { let fixture = Path::new(DATA) .join("fixtures") @@ -71,32 +118,7 @@ fn check_model(tok_file: &str) { if chunk.is_empty() { continue; } - for add_special_tokens in [false, true] { - let expected = released - .encode_fast(chunk, add_special_tokens) - .unwrap() - .get_ids() - .to_vec(); - let got: Vec = pipeline - .encode(chunk, add_special_tokens) - .unwrap() - .iter() - .map(|t| t.id) - .collect(); - if expected != got { - let at = expected - .iter() - .zip(&got) - .position(|(e, g)| e != g) - .unwrap_or(expected.len().min(got.len())); - failures.push(format!( - "{group}/{stem} ({w} B window, add_special_tokens={add_special_tokens}): \ - ids diverge at {at} (expected len {}, got len {})", - expected.len(), - got.len(), - )); - } - } + check(format!("{group}/{stem} ({w} B window)"), chunk); } } assert!( From 0484ea4390833998f2085fcc4149a7c85202f414 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:37:46 +0200 Subject: [PATCH 04/13] perf(pipeline): fuse Metaspace prepend and replace into one normalizer pass Prepend + Replace (or a lone Replace) collapse into a single MetaspaceNormalizer at pipeline build time, so the text is walked once instead of twice. Squashed from perf/fuse-metaspace (no PR). --- tokenizers/tk-encode/examples/normalize_ab.rs | 98 +++++++ .../tk-encode/src/normalizers/metaspace.rs | 258 +++++++++++++++++- .../tk-encode/src/normalizers/replace.rs | 2 +- .../tk-encode/src/pre_tokenizers/metaspace.rs | 12 +- .../tk-encode/src/tokenizer/pipeline.rs | 118 +++++++- tokenizers/tk-encode/tests/pipeline_oracle.rs | 24 ++ 6 files changed, 492 insertions(+), 20 deletions(-) create mode 100644 tokenizers/tk-encode/examples/normalize_ab.rs diff --git a/tokenizers/tk-encode/examples/normalize_ab.rs b/tokenizers/tk-encode/examples/normalize_ab.rs new file mode 100644 index 0000000000..2c8a8304a0 --- /dev/null +++ b/tokenizers/tk-encode/examples/normalize_ab.rs @@ -0,0 +1,98 @@ +//! A/B harness for changes to the pipeline's normalize stage: for each tokenizer config given on +//! the command line, times the `encode_generic` ablation ladder over `data/big.txt` and prints the +//! frame and normalize levels plus the full encode throughput. +//! +//! cargo run --release -p tk-encode --example normalize_ab -- data/llama-2.json data/gemma-4.json +//! +//! Binary layout alone moves numbers by a few percent, so never compare two builds from one +//! process each: build one binary per side, keep both, and alternate runs. + +use std::convert::TryFrom; +use std::hint::black_box; +use std::path::Path; +use std::time::Instant; + +use tk_encode::Tokenizer; +use tk_encode::pipeline::{Model, PipelineTokenizer}; + +const DATA_DIR: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data"); +/// Per-input-overhead is amortized at this size; same regime as `fixture_bench`. +const CHUNK_BYTES: usize = 10 * 1024; +const TOTAL_BYTES: usize = 4 * 1024 * 1024; +const REPS: usize = 9; + +fn median(mut samples: Vec) -> f64 { + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + samples[samples.len() / 2] +} + +/// Median seconds of one warm pass over `chunks` at ladder level `STAGE`; +/// same shape as `fixture_bench::stage_secs`. +fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[&str]) -> f64 { + let mut out = Vec::new(); + let mut pre_tokens = Vec::new(); + let mut scratch = pipeline.get_model().init_scratch(); + let mut run = || { + for chunk in chunks { + out.clear(); + let _ = pipeline.encode_generic::( + chunk, + true, + &mut pre_tokens, + &mut scratch, + &mut out, + ); + black_box(&out); + black_box(&pre_tokens); + } + }; + run(); // warm-up + let mut samples = Vec::with_capacity(REPS); + for _ in 0..REPS { + let start = Instant::now(); + run(); + samples.push(start.elapsed().as_secs_f64()); + } + median(samples) +} + +/// The first `TOTAL_BYTES` of big.txt in `CHUNK_BYTES` pieces, cut on char boundaries. +fn chunks_of(text: &str) -> Vec<&str> { + let mut chunks = Vec::new(); + let mut start = 0; + while start < text.len().min(TOTAL_BYTES) { + let mut end = (start + CHUNK_BYTES).min(text.len()); + while !text.is_char_boundary(end) { + end += 1; + } + chunks.push(&text[start..end]); + start = end; + } + chunks +} + +fn main() { + let text = std::fs::read_to_string(Path::new(DATA_DIR).join("big.txt")) + .expect("data/big.txt (fetch with `make data/big.txt`)"); + let chunks = chunks_of(&text); + let bytes: usize = chunks.iter().map(|c| c.len()).sum(); + + for config in std::env::args().skip(1) { + let tok = Tokenizer::from_file(&config).expect("tokenizer config"); + let pipeline = PipelineTokenizer::try_from(&tok).expect("pipeline builds"); + + let t_frame = stage_secs::<{ PipelineTokenizer::STAGE_FRAME }>(&pipeline, &chunks); + let t_norm = stage_secs::<{ PipelineTokenizer::STAGE_NORMALIZE }>(&pipeline, &chunks); + let t_full = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(&pipeline, &chunks); + + let nspb = |secs: f64| secs * 1e9 / bytes as f64; + println!( + "{config}: frame {:.3} ns/B, normalize {:.3} ns/B (marginal {:.3}), full {:.3} ns/B = {:.1} MB/s", + nspb(t_frame), + nspb(t_norm), + nspb(t_norm - t_frame), + nspb(t_full), + bytes as f64 / t_full / 1e6, + ); + } +} diff --git a/tokenizers/tk-encode/src/normalizers/metaspace.rs b/tokenizers/tk-encode/src/normalizers/metaspace.rs index b1a9486cee..b046fdb70a 100644 --- a/tokenizers/tk-encode/src/normalizers/metaspace.rs +++ b/tokenizers/tk-encode/src/normalizers/metaspace.rs @@ -10,22 +10,45 @@ //! `to_normalizer_and_split`, in [`crate::pre_tokenizers::metaspace`], builds that pair and spells //! out which [`Metaspace`] settings can be rebuilt this way. //! +//! Other configs spell the same rewrite in their `normalizer` field instead: llama-2 as +//! `Prepend("▁")` followed by `Replace(" " -> "▁")`, gemma-4 as the `Replace` alone. Run as +//! declared, every one of those steps writes a full copy of the text. `MetaspaceNormalizer::fuse` +//! recognizes both shapes when the pipeline is built and stands in for them: one pass, one +//! right-sized allocation. +//! //! [`Metaspace`]: crate::pre_tokenizers::metaspace::Metaspace //! [`Split`]: crate::pre_tokenizers::split::Split use std::borrow::Cow; +use crate::normalizers::NormalizerWrapper; +use crate::normalizers::replace::{Replace, ReplacePattern}; use crate::pre_tokenizers::whitespace::WhitespaceSplit; use crate::tokenizer::{Result, pipeline}; +/// When [`MetaspaceNormalizer`] writes a delimiter at the start of the text it is given. +/// +/// The two prepending modes differ only on text that already starts with the delimiter (or with a +/// space the swap turns into one): `IfMissing` skips those, which is how a [`Metaspace`] +/// pre-tokenizer prepends; `Unconditional` marks them again, which is what a `Prepend` normalizer +/// running before the swap does. +/// +/// [`Metaspace`]: crate::pre_tokenizers::metaspace::Metaspace +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) enum PrependMode { + Never, + IfMissing, + Unconditional, +} + /// Writes the delimiter where words start (after a space) #[derive(Debug, Clone, PartialEq)] pub struct MetaspaceNormalizer { /// `▁` (U+2581) for every SentencePiece model we know of, but the config is free to use /// another character. delimiter: char, - /// Write the delimiter at the start of every word, not only the words that followed a space. - prepend: bool, + /// When to write the delimiter at the start of the text. + prepend: PrependMode, /// Throw whitespace away instead of turning it into a delimiter: tabs, newlines and repeated /// spaces leave no trace, and each word keeps only the one delimiter `prepend` writes. This is /// the [`WhitespaceSplit`] that t5 and albert run in front of their `Metaspace`. @@ -33,13 +56,54 @@ pub struct MetaspaceNormalizer { } impl MetaspaceNormalizer { - pub(crate) fn new(delimiter: char, prepend: bool, drop_whitespace: bool) -> Self { + pub(crate) fn new(delimiter: char, prepend: PrependMode, drop_whitespace: bool) -> Self { Self { delimiter, prepend, drop_whitespace, } } + + /// The one-pass stand-in for the leading `steps`, when they spell out this normalizer's job: + /// `Prepend(c)` followed by `Replace(" " -> c)` (llama-2's normalizer), or a `Replace(" " -> c)` + /// on its own (gemma-4's). Returns the stand-in and how many steps it covers; `None` when the + /// steps start with anything else, and the caller keeps them as declared. + /// + /// [`PipelineTokenizer::try_from`] runs the config's normalizer steps through this when the + /// pipeline is built. The config itself is never rewritten. + /// + /// [`PipelineTokenizer::try_from`]: crate::pipeline::PipelineTokenizer + pub(crate) fn fuse(steps: &[NormalizerWrapper]) -> Option<(Self, usize)> { + if let [ + NormalizerWrapper::Prepend(prepend), + NormalizerWrapper::Replace(replace), + .., + ] = steps + && let Some(delimiter) = space_swap(replace) + && prepend.prepend == replace.content + { + return Some((Self::new(delimiter, PrependMode::Unconditional, false), 2)); + } + if let [NormalizerWrapper::Replace(replace), ..] = steps + && let Some(delimiter) = space_swap(replace) + { + return Some((Self::new(delimiter, PrependMode::Never, false), 1)); + } + None + } +} + +/// The delimiter a [`Replace`] swaps every space for: its pattern must be the string `" "` and its +/// content a single char. Anything else (a regex pattern, even one only matching a space; content +/// of any other length) is not checkable as the space swap and returns `None`. +fn space_swap(replace: &Replace) -> Option { + match &replace.pattern { + ReplacePattern::String(pattern) if pattern == " " => {} + _ => return None, + } + let mut chars = replace.content.chars(); + let delimiter = chars.next()?; + chars.next().is_none().then_some(delimiter) } impl pipeline::Normalizer for MetaspaceNormalizer { @@ -48,24 +112,43 @@ impl pipeline::Normalizer for MetaspaceNormalizer { if input.is_empty() { return Ok(Cow::Borrowed(input)); } - // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space - let mut rewritten = String::with_capacity(input.len() + input.len() / 2); if self.drop_whitespace { + // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space + let mut rewritten = String::with_capacity(input.len() + input.len() / 2); // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the // words back one after the other, each with its own delimiter. let mut words = Vec::new(); pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?; for span in &words { let word = &input[span.range()]; - // The text may already hold delimiters of its own — never write a second one. - if self.prepend && !word.starts_with(self.delimiter) { + let prepend = match self.prepend { + PrependMode::Never => false, + // The text may already hold delimiters of its own: never write a second one. + PrependMode::IfMissing => !word.starts_with(self.delimiter), + PrependMode::Unconditional => true, + }; + if prepend { rewritten.push(self.delimiter); } rewritten.push_str(word); } + Ok(Cow::Owned(rewritten)) } else { - // Prepend the delimiter if self.prepend is true - if self.prepend && !input.starts_with(' ') && !input.starts_with(self.delimiter) { + let prepend = match self.prepend { + PrependMode::Never => false, + // A leading space counts as already marked: the swap turns it into a delimiter. + PrependMode::IfMissing => { + !input.starts_with(' ') && !input.starts_with(self.delimiter) + } + PrependMode::Unconditional => true, + }; + // Nothing to prepend and nothing to swap: hand the input back instead of copying it. + if !prepend && memchr::memchr(b' ', input.as_bytes()).is_none() { + return Ok(Cow::Borrowed(input)); + } + // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space + let mut rewritten = String::with_capacity(input.len() + input.len() / 2); + if prepend { rewritten.push(self.delimiter); } // Only spaces become delimiters; tabs and newlines are left alone @@ -76,7 +159,162 @@ impl pipeline::Normalizer for MetaspaceNormalizer { rest = &rest[space + 1..]; } rewritten.push_str(rest); + Ok(Cow::Owned(rewritten)) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tokenizer::pipeline::normalize_all; + + fn step(json: &str) -> NormalizerWrapper { + serde_json::from_str(json).unwrap() + } + + /// llama-2's two normalizer steps. + const PREPEND: &str = r#"{"type":"Prepend","prepend":"▁"}"#; + /// gemma-4's whole normalizer, and llama-2's second step. + const SWAP: &str = r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#; + + /// Every chunk shape the rewrite can meet: empty, already marked, leading, trailing and + /// repeated spaces, no spaces at all, other whitespace. + const TEXTS: &[&str] = &[ + "", + "hello world", + "hello world", + " leading", + "trailing ", + " both ", + "▁already marked", + "▁", + "no_spaces", + "one\ttab\nand a newline", + " ", + "a▁b c", + ]; + + /// The stand-in must rewrite every text exactly as the steps it covers would. + fn assert_fused_matches_steps(jsons: &[&str]) { + let steps: Vec = jsons.iter().map(|json| step(json)).collect(); + let (fused, covered) = MetaspaceNormalizer::fuse(&steps).expect("this shape fuses"); + assert_eq!(covered, steps.len()); + for text in TEXTS { + let expected = normalize_all(&steps, text).unwrap(); + let got = pipeline::Normalizer::normalize(&fused, text).unwrap(); + assert_eq!(got, expected, "{text:?}"); } - Ok(Cow::Owned(rewritten)) + } + + #[test] + fn fused_prepend_and_swap_match_the_two_steps() { + assert_fused_matches_steps(&[PREPEND, SWAP]); + } + + #[test] + fn fused_swap_matches_the_lone_step() { + assert_fused_matches_steps(&[SWAP]); + } + + /// Nothing ties the delimiter to `▁`; any single char must fuse the same way. + #[test] + fn fused_ascii_delimiter_matches_its_steps() { + assert_fused_matches_steps(&[ + r#"{"type":"Prepend","prepend":"_"}"#, + r#"{"type":"Replace","pattern":{"String":" "},"content":"_"}"#, + ]); + } + + /// The pair stands in for a `Prepend`, which marks a chunk even when it already starts with + /// the delimiter; the lone swap never prepends. `covered` tells the caller how many steps to + /// skip. + #[test] + fn the_pair_prepends_unconditionally_and_the_lone_swap_never_does() { + let steps = [step(PREPEND), step(SWAP)]; + let (fused, covered) = MetaspaceNormalizer::fuse(&steps).unwrap(); + assert_eq!(covered, 2); + assert_eq!( + fused, + MetaspaceNormalizer::new('▁', PrependMode::Unconditional, false) + ); + let (fused, covered) = MetaspaceNormalizer::fuse(&steps[1..]).unwrap(); + assert_eq!(covered, 1); + assert_eq!( + fused, + MetaspaceNormalizer::new('▁', PrependMode::Never, false) + ); + } + + /// A chunk with nothing to swap is handed back without a copy, as `Replace` hands it back. + #[test] + fn a_swap_with_no_spaces_borrows() { + let steps = [step(SWAP)]; + let (fused, _) = MetaspaceNormalizer::fuse(&steps).unwrap(); + assert!(matches!( + pipeline::Normalizer::normalize(&fused, "no_spaces").unwrap(), + Cow::Borrowed(_) + )); + } + + #[test] + fn refuses_a_replace_that_is_not_the_space_swap() { + let refused = [ + // Swaps something else entirely; running it as the space swap would corrupt the text. + ( + "another literal", + r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#, + ), + // The delimiter is one char; two of them per space is not this normalizer's rewrite. + ( + "multi-char content", + r#"{"type":"Replace","pattern":{"String":" "},"content":"▁▁"}"#, + ), + // Deletes spaces instead of marking them. + ( + "empty content", + r#"{"type":"Replace","pattern":{"String":" "},"content":""}"#, + ), + // Only single spaces are swapped one-for-one. + ( + "a two-space pattern", + r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#, + ), + ]; + for (name, json) in refused { + assert!(MetaspaceNormalizer::fuse(&[step(json)]).is_none(), "{name}"); + } + } + + /// A regex spelling is refused even when it happens to match only a space: fusing checks the + /// pattern structurally and does not interpret regex syntax. + #[cfg(feature = "fancy-regex")] + #[test] + fn refuses_a_regex_pattern_even_one_matching_a_space() { + let json = r#"{"type":"Replace","pattern":{"Regex":" "},"content":"▁"}"#; + assert!(MetaspaceNormalizer::fuse(&[step(json)]).is_none()); + } + + /// A `Prepend` writing anything but the swap's delimiter is a different rewrite; the pair must + /// not fuse. The swap after it still can, on its own, once the caller reaches it. + #[test] + fn refuses_a_prepend_of_something_else() { + for prepend in [ + r#"{"type":"Prepend","prepend":"x"}"#, + r#"{"type":"Prepend","prepend":"▁▁"}"#, + ] { + let steps = [step(prepend), step(SWAP)]; + assert!(MetaspaceNormalizer::fuse(&steps).is_none(), "{prepend}"); + assert!(MetaspaceNormalizer::fuse(&steps[1..]).is_some()); + } + } + + /// The pair only fuses when the two steps are adjacent: another step in between sees the text + /// mid-rewrite, and fusing across it would change what that step sees. + #[test] + fn refuses_a_separated_pair() { + let steps = [step(PREPEND), step(r#"{"type":"Lowercase"}"#), step(SWAP)]; + assert!(MetaspaceNormalizer::fuse(&steps).is_none()); + assert!(MetaspaceNormalizer::fuse(&steps[2..]).is_some()); } } diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index 03e207360d..5f2ea686ce 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -70,7 +70,7 @@ impl Search { #[derive(Debug, Serialize, Deserialize)] #[serde(tag = "type", try_from = "ReplaceDeserializer")] pub struct Replace { - pattern: ReplacePattern, + pub(crate) pattern: ReplacePattern, pub content: String, #[serde(skip)] search: Search, diff --git a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs index 3d940b4750..9afd1d1724 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/metaspace.rs @@ -1,4 +1,4 @@ -use crate::normalizers::metaspace::MetaspaceNormalizer; +use crate::normalizers::metaspace::{MetaspaceNormalizer, PrependMode}; use crate::pre_tokenizers::PreTokenizerWrapper; use crate::pre_tokenizers::split::Split; use crate::tokenizer::{Decoder, PreTokenizedString, PreTokenizer, Result, SplitDelimiterBehavior}; @@ -212,15 +212,17 @@ fn normalizer_and_split( return None; } let prepend = match metaspace.prepend_scheme { - PrependScheme::Always => true, - PrependScheme::Never => false, + // `Metaspace` swaps spaces first and then marks only text not already starting with the + // delimiter, so its `always` scheme is the conditional prepend. + PrependScheme::Always => PrependMode::IfMissing, + PrependScheme::Never => PrependMode::Never, // `First` writes the delimiter only on the piece at the very start of the text it came from. // A normalizer is handed one chunk at a time, without that context. PrependScheme::First => return None, }; // Removes whitespaces and does not prepend words: nothing would show where words begin - // The output is one big continuous blob of words hlued together - if drop_whitespace && !prepend { + // The output is one big continuous blob of words glued together + if drop_whitespace && prepend == PrependMode::Never { return None; } let delimiter = metaspace.replacement; diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 4dc3d5e285..29aefd51d1 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -108,9 +108,11 @@ pub(crate) fn normalize_all<'a, N: Normalizer>( #[allow(clippy::large_enum_variant)] #[derive(Debug)] enum PipelineNormalizer { - /// The `normalizer` field of the config, as-is. + /// One step of the config's `normalizer` field, as-is (a `Sequence`'s members each get their + /// own entry). Declared(NormalizerWrapper), - /// The text-rewriting half of a `Metaspace` pre-tokenizer. + /// The one-pass stand-in for declared steps that spell out the SentencePiece space rewrite + /// ([`MetaspaceNormalizer::fuse`]), or the text-rewriting half of a `Metaspace` pre-tokenizer. Metaspace(MetaspaceNormalizer), } @@ -539,7 +541,29 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { fn try_from(tok: &Tokenizer) -> Result { let mut normalizers = Vec::new(); if let Some(declared) = tok.get_normalizer() { - normalizers.push(PipelineNormalizer::Declared(declared.clone())); + // A `Sequence` runs its members in order, exactly what this list does, so its members + // join the list one by one. That puts them in [`MetaspaceNormalizer::fuse`]'s reach: + // steps spelling out the SentencePiece space rewrite (llama-2 as two whole-chunk + // copies, `Prepend("▁")` then `Replace(" " -> "▁")`; gemma-4 as the `Replace` alone) + // are swapped for the normalizer doing that job in one pass. Only this in-memory + // pipeline is rewritten; the config still serializes as declared. + let steps: &[NormalizerWrapper] = match declared { + NormalizerWrapper::Sequence(sequence) => sequence.as_ref(), + single => std::slice::from_ref(single), + }; + let mut index = 0; + while index < steps.len() { + match MetaspaceNormalizer::fuse(&steps[index..]) { + Some((fused, covered)) => { + normalizers.push(PipelineNormalizer::Metaspace(fused)); + index += covered; + } + None => { + normalizers.push(PipelineNormalizer::Declared(steps[index].clone())); + index += 1; + } + } + } } // A `Metaspace` pre-tokenizer does two jobs at once: it writes `▁` delimiters into the text, @@ -1226,11 +1250,97 @@ mod tests { .iter() .map(|t| t.id) .collect(); - // Not the unk id: both the `Replace` and the `Split` really ran on the literal path. + // Not the unk id: the string-pattern config ran end-to-end with no regex backend (the + // `Replace` fuses into the Metaspace normalizer; the `Split` runs on the literal path). assert_eq!(ids, [1, 2]); assert_pipeline_matches_reference(&tok, "hello world"); } + /// llama-2's declared normalizer is `Prepend("▁")` then `Replace(" " -> "▁")`, two steps that + /// each copy the chunk. Building the pipeline must swap them for the one-pass + /// [`MetaspaceNormalizer`] (see [`MetaspaceNormalizer::fuse`]), and must not touch the source + /// config while doing so. + #[test] + fn try_from_fuses_the_declared_space_rewrite() { + let normalizer: NormalizerWrapper = serde_json::from_str( + r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#, + ) + .unwrap(); + let mut tok = wordlevel_tokenizer(vec![("", 0)], None); + tok.with_normalizer(Some(normalizer)).unwrap(); + + let declared = serde_json::to_string(tok.get_normalizer().unwrap()).unwrap(); + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!( + pipeline.normalizers.as_slice(), + [PipelineNormalizer::Metaspace(_)] + ), + "{:?}", + pipeline.normalizers + ); + assert_eq!( + serde_json::to_string(tok.get_normalizer().unwrap()).unwrap(), + declared, + ); + } + + /// Normalization runs on each chunk between special tokens, so the fused prepend must too: + /// the text after `` starts with a space, and `Prepend` + swap turn that into two + /// delimiters. Only `"▁▁world"` is in the vocabulary; a prepend skipping the marked chunk + /// would produce the unk id instead. + #[test] + fn fused_metaspace_normalizer_prepends_per_chunk() { + let normalizer: NormalizerWrapper = serde_json::from_str( + r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#, + ) + .unwrap(); + let mut tok = wordlevel_tokenizer(vec![("", 0), ("▁hello▁", 1), ("▁▁world", 2)], None); + tok.with_normalizer(Some(normalizer)).unwrap(); + tok.with_pre_tokenizer(None::); + let _ = tok.add_special_tokens([crate::AddedToken::from("", true)]); + + let ids: Vec = PipelineTokenizer::try_from(&tok) + .unwrap() + .encode("hello world", false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(ids, [1, 3, 2]); + assert_pipeline_matches_reference(&tok, "hello world"); + } + + /// A `Replace` swapping anything but a single space stays a declared step, and the ids stay + /// those of the declared rewrite. + #[test] + fn a_replace_that_is_not_the_space_swap_is_left_declared() { + let normalizer: NormalizerWrapper = + serde_json::from_str(r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#) + .unwrap(); + let mut tok = wordlevel_tokenizer(vec![("", 0), ("y a", 1)], None); + tok.with_normalizer(Some(normalizer)).unwrap(); + tok.with_pre_tokenizer(None::); + + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!( + pipeline.normalizers.as_slice(), + [PipelineNormalizer::Declared(_)] + ), + "{:?}", + pipeline.normalizers + ); + let ids: Vec = pipeline + .encode("x a", false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(ids, [1]); + assert_pipeline_matches_reference(&tok, "x a"); + } + #[test] fn segment_iterator_yields_text_and_specials_in_order() { let input = "aabbcc"; diff --git a/tokenizers/tk-encode/tests/pipeline_oracle.rs b/tokenizers/tk-encode/tests/pipeline_oracle.rs index 1582c2f98c..6befe4ae01 100644 --- a/tokenizers/tk-encode/tests/pipeline_oracle.rs +++ b/tokenizers/tk-encode/tests/pipeline_oracle.rs @@ -40,6 +40,26 @@ const HOSTILE: &[&str] = &[ " leading and trailing ", ]; +/// Chunk shapes the SentencePiece space rewrite gets wrong first: empty input, text already +/// starting with the `▁` delimiter, leading, trailing and repeated spaces, other whitespace, and +/// special tokens cutting the text into chunks (normalization, and with it the delimiter prepend, +/// runs on each chunk). Strings that are not special tokens for a given model just encode as +/// text, so the whole list is checked for every model. +const EDGE_TEXTS: &[&str] = &[ + "", + " leading", + "trailing ", + " both ", + "a b c", + "▁already marked", + "▁", + "no_spaces", + "\ttab\nand newline", + "hello world", + "hugging face", + "a b c ", +]; + fn check_model(tok_file: &str) { let path = Path::new(DATA).join(tok_file); // The legacy `Tokenizer` only *builds* the pipeline (its sole constructor @@ -99,6 +119,10 @@ fn check_model(tok_file: &str) { check(format!("hostile/{i}"), chunk); } + for &text in EDGE_TEXTS { + check(format!("edge {text:?}"), text); + } + for &(group, stem) in FIXTURES { let fixture = Path::new(DATA) .join("fixtures") From 820d436f2c6f53d78b0615caf3245d27ec53e95a Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:24:57 +0200 Subject: [PATCH 05/13] test(pipeline): pin that the fused normalizer and proven cuts both fire The fuse (perf/fuse-metaspace) and the cuts (perf/metaspace-proven-cuts) were written on branches that did not know about each other. Both detections read the declared config, so they compose, but nothing asserted that: an encode oracle failure would only show if one of them also changed ids. This pins the pipeline shape for the real llama-2 and gemma-4 configs. Co-Authored-By: Claude Fable 5 --- .../tk-encode/src/tokenizer/pipeline.rs | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 29aefd51d1..22fbc84c89 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1285,6 +1285,36 @@ mod tests { ); } + /// The real llama-2 and gemma-4 files qualify for two independent rewrites: their declared + /// space rewrite fuses into the one-pass [`MetaspaceNormalizer`], and the whole-chunk text + /// they hand the model is cut into words by [`ProvenCuts`]. Both read the declared config, + /// not each other's output, so building the pipeline must apply both; this pins that neither + /// detection hides the other. Skipped when the files are not fetched. + #[test] + fn real_sentencepiece_configs_get_the_fused_normalizer_and_proven_cuts() { + for file in ["llama-2.json", "gemma-4.json"] { + let path = format!("../data/{file}"); + if !std::path::Path::new(&path).exists() { + eprintln!("skip {file}: not present (fetch with `make bench-models`)"); + continue; + } + let tok = crate::Tokenizer::from_file(&path).unwrap(); + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + matches!( + pipeline.normalizers.as_slice(), + [PipelineNormalizer::Metaspace(_)] + ), + "{file}: the declared space rewrite should fuse, got {:?}", + pipeline.normalizers, + ); + assert!( + matches!(pipeline.pre_tokenizer, PipelinePreTokenizer::ProvenCuts(_)), + "{file}: the whole-chunk text should get proven cuts", + ); + } + } + /// Normalization runs on each chunk between special tokens, so the fused prepend must too: /// the text after `` starts with a space, and `Prepend` + swap turn that into two /// delimiters. Only `"▁▁world"` is in the vocabulary; a prepend skipping the marked chunk From 8cb43d2c1407b5997948926ded598670073ef252 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:09:39 +0200 Subject: [PATCH 06/13] perf(bpe): seed byte-level words with proven whole-character tokens Ports the char-fold idea from Arthur's tables engine (#2241): at build time, prove from the vocabulary which multi-byte characters always assemble into one token (lowest-rank replay, no boundary neighbour able to pre-empt a step), and seed the merge loop with that token instead of 2-3 byte symbols. Differences from #2241: external ids throughout (no internal renumbering), seeding feeds the existing exact merge loop instead of the WIP multipass engine, and the WordCache above it is untouched. The pair table / grid parts stay behind until that engine is green. llama-3 folds 298 characters (185 CJK); model stage on the jpn fixture, cache off: 37.0 -> 29.3 ns/B (-21%). Ids pinned equal with the fold blanked (llama-3 A/B over CJK/Greek/Cyrillic/emoji corpora) and by the 9/9 encode oracle; thief and non-BMP cases covered in unit tests. Co-Authored-By: Claude Fable 5 --- .../src/models/bpe/bytelevel_folding.rs | 327 ++++++++++++++++++ tokenizers/tk-encode/src/models/bpe/mod.rs | 1 + tokenizers/tk-encode/src/models/bpe/model.rs | 190 +++++++++- 3 files changed, 513 insertions(+), 5 deletions(-) create mode 100644 tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs diff --git a/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs b/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs new file mode 100644 index 0000000000..a19a137129 --- /dev/null +++ b/tokenizers/tk-encode/src/models/bpe/bytelevel_folding.rs @@ -0,0 +1,327 @@ +//! Which characters a byte-level vocabulary can emit as one token instead of as their bytes. +//! +//! A byte-level model's atoms are the 256 bytes, so a multi-byte character like え reaches the +//! merge loop as three symbols that the merges then reassemble. When that reassembly is +//! predetermined, seeding the merge loop with the character's own token skips the work. Two +//! conditions make it predetermined: +//! +//! 1. The character's bytes must collapse into exactly one symbol when the merges are replayed +//! the way BPE picks them: lowest rank first. +//! 2. No step of that replay may be taken over from outside the character. The merge loop does +//! not know where the character ends: if a neighbouring symbol can merge with the character's +//! first or last symbol at a lower rank, that merge fires first and the assembly never +//! happens. +//! +//! A character failing either test gets no table entry; its bytes go through the merge loop as +//! usual, which is always exact. The fold is a shortcut, never a requirement. + +use std::cmp; + +use crate::models::bpe::MergeMap; +use crate::utils::byte_level::CHAR_BYTES_LOOKUP; + +/// A character outside the table, or outside the Basic Multilingual Plane. The seeding loop +/// falls back to per-byte symbols when it reads this. +pub(super) const NO_FOLD: u32 = u32::MAX; + +/// `table[codepoint]` is the id of the token this character folds to, or [`NO_FOLD`]. +/// +/// Indexed by Basic Multilingual Plane codepoints only: a `char` above `u16::MAX` never folds. +/// Single-byte characters are also left out, their byte symbol already is the seed. +pub(super) type FoldTable = Box<[u32; FOLD_TABLE_LEN]>; +pub(super) const FOLD_TABLE_LEN: usize = 1 << 16; + +/// Builds the fold table for a byte-level vocabulary. +/// +/// `vocab` still spells its tokens in byte-level characters ("é" for é), which is what +/// [`ByteLevelFold::fold`] expects; call this before the store is rebuilt on raw bytes. +pub(super) fn build_fold_table(vocab: &[(String, u32)], merges: &MergeMap) -> FoldTable { + let fold = ByteLevelFold::new(vocab, merges); + let mut table = vec![NO_FOLD; FOLD_TABLE_LEN]; + for (token, id) in vocab { + if let Fold::Folds(ch, id) = fold.fold(token, *id) + && ch.len_utf8() > 1 + && (ch as usize) < FOLD_TABLE_LEN + { + table[ch as usize] = id; + } + } + table + .into_boxed_slice() + .try_into() + .expect("length is FOLD_TABLE_LEN") +} + +/// What one vocab token is worth to the fold table. +pub(super) enum Fold { + /// A single character whose bytes assemble to exactly this token, un-stealably. + Folds(char, u32), + /// Formable, but some step could be taken over by a neighbour. + Unsafe, + /// Not a single character, or its bytes never assemble at all. Nothing to record. + Skip, +} + +pub(super) struct ByteLevelFold<'a> { + /// byte -> id of that byte's own one-character token. A byte's value is not its id + /// (gpt2: 0x41 -> 32, 0x20 -> 220), which is why this indirection exists. + byte_token: [u32; 256], + /// `stolen_from_left[id]`: the lowest rank at which a left neighbour merges with `id`, + /// counting only neighbours reachable at a character boundary. Same on the other side for + /// `stolen_from_right`. + stolen_from_left: Vec, + stolen_from_right: Vec, + merges: &'a MergeMap, +} + +impl<'a> ByteLevelFold<'a> { + pub(super) fn new(vocab: &[(String, u32)], merges: &'a MergeMap) -> Self { + let mut byte_token = [u32::MAX; 256]; + for (token, id) in vocab { + let mut chars = token.chars(); + if let (Some(ch), None) = (chars.next(), chars.next()) + && let Some(&b) = CHAR_BYTES_LOOKUP.get(&ch) + { + byte_token[b as usize] = *id; + } + } + + let (stolen_from_left, stolen_from_right) = boundary_merge_ranks(vocab, merges); + + Self { + byte_token, + stolen_from_left, + stolen_from_right, + merges, + } + } + + /// Verdict for `token`, whose id is `id`. + pub(super) fn fold(&self, token: &str, id: u32) -> Fold { + let Some(bytes) = token + .chars() + .map(|ch| CHAR_BYTES_LOOKUP.get(&ch).copied()) + .collect::>>() + else { + // A character with no byte-level mapping: an added token, not text. + return Fold::Skip; + }; + // The table is keyed by codepoint, so only single-character tokens can go in it. This + // also drops lone bytes >= 0x80, which are not characters on their own. + let Ok(text) = std::str::from_utf8(&bytes) else { + return Fold::Skip; + }; + let mut it = text.chars(); + let (Some(ch), None) = (it.next(), it.next()) else { + return Fold::Skip; + }; + + let mut running: Vec = bytes.iter().map(|&b| self.byte_token[b as usize]).collect(); + if running.contains(&u32::MAX) { + return Fold::Skip; // a byte with no token of its own: never assemblable + } + // Replay the assembly the way the merge loop picks: the lowest rank among the pairs + // still standing, until one symbol is left. + while running.len() > 1 { + let mut best: Option<(usize, u32, u32)> = None; + for i in 0..running.len() - 1 { + let pair = (running[i], running[i + 1]); + if let Some((rank, product)) = self.merges.get(&pair) + && best.is_none_or(|(_, best_rank, _)| *rank < best_rank) + { + best = Some((i, *rank, *product)); + } + } + let Some((i, rank, product)) = best else { + return Fold::Skip; // stuck above one symbol: reference BPE stops here too + }; + if rank >= self.stolen_from_left[running[0] as usize] + || rank >= self.stolen_from_right[*running.last().unwrap() as usize] + { + return Fold::Unsafe; + } + running[i] = product; + running.remove(i + 1); + } + + debug_assert_eq!(running[0], id); + Fold::Folds(ch, running[0]) + } +} + +/// A folded character's edges are character boundaries by construction, and UTF-8 +/// pins down what may be there: +/// +/// - right neighbour = the next character's FIRST byte -> ASCII or a lead byte, never 0x80..=0xBF +/// - left neighbour = the previous character's LAST byte -> never a lead byte, so always < 0xC0 +fn boundary_merge_ranks(vocab: &[(String, u32)], merges: &MergeMap) -> (Vec, Vec) { + let max_id = vocab + .iter() + .map(|(_, id)| *id) + .chain(merges.iter().flat_map(|((a, b), (_, id))| [*a, *b, *id])) + .max() + .map_or(0, |id| id as usize + 1); + + // First and last real byte of every token. 0xFF marks a token with no byte spelling (an + // added token): it counts as a possible stealer on the right, and as none on the left, + // which errs towards refusing a fold. + let (mut first, mut last) = (vec![0xFFu8; max_id], vec![0xFFu8; max_id]); + for (token, id) in vocab { + let Some(bytes) = token + .chars() + .map(|c| CHAR_BYTES_LOOKUP.get(&c).copied()) + .collect::>>() + else { + continue; + }; + if let (Some(f), Some(l)) = (bytes.first(), bytes.last()) { + first[*id as usize] = *f; + last[*id as usize] = *l; + } + } + // 0xC0/0xC1 are overlong lead bytes and cannot occur on either side. + let starts_at_boundary = |id: u32| first[id as usize] < 0x80 || first[id as usize] >= 0xC2; + let ends_at_boundary = |id: u32| last[id as usize] < 0xC0; + + let mut stolen_from_left = vec![u32::MAX; max_id]; + let mut stolen_from_right = vec![u32::MAX; max_id]; + for ((a, b), (rank, _)) in merges.iter() { + if *a as usize >= max_id || *b as usize >= max_id { + continue; + } + if starts_at_boundary(*b) { + stolen_from_right[*a as usize] = cmp::min(stolen_from_right[*a as usize], *rank); + } + if ends_at_boundary(*a) { + stolen_from_left[*b as usize] = cmp::min(stolen_from_left[*b as usize], *rank); + } + } + (stolen_from_left, stolen_from_right) +} + +#[cfg(test)] +mod test { + use super::{ByteLevelFold, Fold, NO_FOLD, build_fold_table}; + use crate::models::bpe::MergeMap; + use crate::utils::byte_level::BYTES_CHAR_LOOKUP; + + /// 'é' is U+00E9 = bytes C3 A9; both are printable latin-1, so the byte-level names are the + /// identity chars 'Ã' and '©' and the vocab spells the character "é". + fn setup(extra_merge: bool) -> (Vec<(String, u32)>, MergeMap) { + let vocab = Vec::from([ + ("Ã".to_string(), 0), // byte 0xC3 + ("©".to_string(), 1), // byte 0xA9 + ("é".to_string(), 2), // the character é + ("x".to_string(), 3), + ("xÃ".to_string(), 4), + ]); + let mut merges = MergeMap::new(); + merges.insert((0, 1), (1, 2)); // à + © -> é at rank 1 + if extra_merge { + // x + à at rank 0: a left neighbour "x" grabs our first byte first, so the + // assembly of é never happens and folding it would be wrong. + merges.insert((3, 0), (0, 4)); + } + (vocab, merges) + } + + #[test] + fn folds_when_nothing_can_steal_an_edge() { + let (vocab, merges) = setup(false); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("é", 2), Fold::Folds('é', 2))); + } + + #[test] + fn rejects_a_boundary_steal() { + let (vocab, merges) = setup(true); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("é", 2), Fold::Unsafe)); + } + + #[test] + fn skips_what_is_not_one_character() { + let (vocab, merges) = setup(false); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("xÃ", 4), Fold::Skip)); // two characters once decoded + assert!(matches!(f.fold("<|endoftext|>", 9), Fold::Skip)); // '<' is fine, '|' is not remapped + assert!(matches!(f.fold("Ã", 0), Fold::Skip)); // lone 0xC3 is not valid UTF-8 + assert!(matches!(f.fold("x", 3), Fold::Folds('x', 3))); // ASCII needs no assembly + } + + #[test] + fn table_keeps_the_fold_and_drops_ascii() { + let (vocab, merges) = setup(false); + let table = build_fold_table(&vocab, &merges); + assert_eq!(table['é' as usize], 2); + // 'x' folds but single-byte characters stay out: their byte symbol already is the seed. + assert_eq!(table['x' as usize], NO_FOLD); + } + + // Byte-level merges in a gpt2-like encoding. Byte level rewrites the vocab so every byte is a + // printable char; bytes 0x80..=0xA0 become U+0122.. and 0xAE..=0xFF stay themselves. + // U+671D 朝 → E6 9C 9D -> 'æ','ľ','Ŀ' + // U+65E5 日 → E6 97 A5 -> 'æ','Ĺ','¥' + // + // 朝 assembles with ('æ','ľ') and ('æľ','Ŀ'), so it may only merge if no merge pair can take + // an edge symbol first. We add such a pair: ('Ŀ','æ') 9D E6, which appears in 朝朝 and 朝日 + // at the boundary `.. 9D | E6 ..`. It has to be a LEAD byte (E6) doing the stealing: the symbol + // after a complete character is always the next character's first byte. + enum Thief { + None, + Lead, + Continuation, + } + + fn cjk_vocab(thief: Thief) -> (Vec<(String, u32)>, MergeMap) { + assert_eq!( + [0xE6u8, 0x9C, 0x9D].map(|b| BYTES_CHAR_LOOKUP[b as usize]), + ['æ', 'ľ', 'Ŀ'] + ); + let mut vocab = Vec::from([ + ("æ".to_string(), 0), // E6 + ("ľ".to_string(), 1), // 9C + ("Ŀ".to_string(), 2), // 9D + ("æľ".to_string(), 3), // E6 9C + ("æľĿ".to_string(), 4), // E6 9C 9D = 朝 + ]); + let mut merges = MergeMap::new(); + // (left, right) -> (rank, product). Ranks leave room below for the thief. + merges.insert((0, 1), (1, 3)); // 'æ' + 'ľ' -> "æľ" + merges.insert((3, 2), (2, 4)); // "æľ" + 'Ŀ' -> 朝 + match thief { + Thief::None => {} + Thief::Lead => { + vocab.push(("Ŀæ".to_string(), 5)); // 9D E6, straddles a character boundary + merges.insert((2, 0), (0, 5)); // rank 0, below every step of 朝's assembly + } + Thief::Continuation => { + vocab.push(("Ģ".to_string(), 5)); // 80 + vocab.push(("ĿĢ".to_string(), 6)); // 9D 80, never at a boundary + merges.insert((2, 5), (0, 6)); + } + } + (vocab, merges) + } + + #[test] + fn folds_a_cjk_char_when_no_neighbour_can_steal() { + let (vocab, merges) = cjk_vocab(Thief::None); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("æľĿ", 4), Fold::Folds('朝', 4))); + } + + #[test] + fn refuses_the_same_char_once_a_lead_byte_can_steal() { + let (vocab, merges) = cjk_vocab(Thief::Lead); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("æľĿ", 4), Fold::Unsafe)); + } + + #[test] + fn a_continuation_byte_cannot_steal_so_it_still_folds() { + let (vocab, merges) = cjk_vocab(Thief::Continuation); + let f = ByteLevelFold::new(&vocab, &merges); + assert!(matches!(f.fold("æľĿ", 4), Fold::Folds('朝', 4))); + } +} diff --git a/tokenizers/tk-encode/src/models/bpe/mod.rs b/tokenizers/tk-encode/src/models/bpe/mod.rs index 6e1cb2da93..06a7ebc5af 100644 --- a/tokenizers/tk-encode/src/models/bpe/mod.rs +++ b/tokenizers/tk-encode/src/models/bpe/mod.rs @@ -1,6 +1,7 @@ //! [Byte Pair Encoding](https://www.aclweb.org/anthology/P16-1162/) model. use std::{iter, mem}; +mod bytelevel_folding; mod model; mod serialization; pub mod word; diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index dbcfb254c1..1ab2d2ae78 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -1,4 +1,4 @@ -use super::{super::OrderedVocabIter, Error, Pair, Word}; +use super::{super::OrderedVocabIter, Error, Pair, Word, bytelevel_folding}; use crate::models::bpe::Merge; use crate::pipeline::{self, ModelScratch, PipelineToken}; use crate::tokenizer::{Model, Result, Token}; @@ -688,6 +688,9 @@ pub struct PipelineBPE { enum Atoms { Bytes { byte_to_id: [u32; 256], + /// Multi-byte characters proven to always assemble into one token; see + /// [`bytelevel_folding`]. The merge loop is seeded with that token directly. + fold: bytelevel_folding::FoldTable, }, Chars { byte_fallback: Option<[u32; 256]>, @@ -718,6 +721,7 @@ impl PipelineBPE { } = model; let (vocab, atoms) = if with_byte_level { + let fold = bytelevel_folding::build_fold_table(&vocab.content(), &merges); let mut vocab = BucketVocabStore::build(vocab.byte_content()); vocab = byte_level::transform_vocab(vocab); let mut byte_to_id = [0u32; 256]; @@ -726,7 +730,7 @@ impl PipelineBPE { .get_bytes(&[b]) .ok_or(Error::ByteAtomOutOfVocabulary(b))?; } - (vocab, Atoms::Bytes { byte_to_id }) + (vocab, Atoms::Bytes { byte_to_id, fold }) } else { let vocab = BucketVocabStore::build(vocab.byte_content()); let unk_token = if let Some(unk_str) = unk_token { @@ -787,9 +791,32 @@ impl PipelineBPE { ) { word.clear(); match &self.atoms { - Atoms::Bytes { byte_to_id } => { - for &b in sequence.as_bytes() { - word.add(byte_to_id[b as usize], 1); + Atoms::Bytes { byte_to_id, fold } => { + let bytes = sequence.as_bytes(); + let mut i = 0; + while i < bytes.len() { + let b = bytes[i]; + if b < 0x80 { + word.add(byte_to_id[b as usize], 1); + i += 1; + continue; + } + // A multi-byte character seeds as its own token when the fold proved the + // merges would assemble it anyway, and as its bytes otherwise. + let ch = sequence[i..].chars().next().unwrap(); + let len = ch.len_utf8(); + let id = match (ch as usize) < bytelevel_folding::FOLD_TABLE_LEN { + true => fold[ch as usize], + false => bytelevel_folding::NO_FOLD, + }; + if id != bytelevel_folding::NO_FOLD { + word.add(id, len); + } else { + for &b in &bytes[i..i + len] { + word.add(byte_to_id[b as usize], 1); + } + } + i += len; } } Atoms::Chars { @@ -1754,6 +1781,159 @@ mod tests { ); } + /// Blanks a model's fold table, so it seeds every character as bytes again. + fn without_fold(mut model: PipelineBPE) -> PipelineBPE { + let Atoms::Bytes { fold, .. } = &mut model.atoms else { + panic!("not a byte-level model"); + }; + *fold = vec![bytelevel_folding::NO_FOLD; bytelevel_folding::FOLD_TABLE_LEN] + .into_boxed_slice() + .try_into() + .unwrap(); + model + } + + /// 'é' (bytes C3 A9) has a token and an unstealable assembly, so the fold seeds it as + /// that token; a thief merge on its first byte must blank the entry, and a character + /// outside the Basic Multilingual Plane never enters the table. Ids stay the same + /// either way. + #[test] + fn folded_char_seeds_as_its_token() { + let n = |b: u8| BYTES_CHAR_LOOKUP[b as usize].to_string(); + let mut vocab: Vocab = (0..=255u8).map(|b| (n(b), u32::from(b))).collect(); + vocab.insert(n(0xC3) + &n(0xA9), 300); + let merges = vec![(n(0xC3), n(0xA9))]; + let bpe = BpeBuilder::default() + .vocab_and_merges(vocab.clone(), merges.clone()) + .build() + .unwrap(); + let pipeline = PipelineBPE::from_bpe(bpe, true).unwrap(); + + let Atoms::Bytes { fold, .. } = &pipeline.atoms else { + panic!("not a byte-level model"); + }; + assert_eq!(fold['é' as usize], 300); + assert_eq!(pipeline_ids(&pipeline, "é"), vec![300]); + // 4-byte characters stay out of the table and encode as their bytes. + assert_eq!(pipeline_ids(&pipeline, "🎉").len(), 4); + + // "x" + C3 at rank 0 fires before é's assembly, so é may not fold; the byte + // path still merges it wherever no thief is adjacent. + vocab.insert("x".to_string() + &n(0xC3), 301); + let mut merges = merges; + merges.insert(0, ("x".to_string(), n(0xC3))); + let bpe = BpeBuilder::default() + .vocab_and_merges(vocab, merges) + .build() + .unwrap(); + let pipeline = PipelineBPE::from_bpe(bpe, true).unwrap(); + let Atoms::Bytes { fold, .. } = &pipeline.atoms else { + panic!("not a byte-level model"); + }; + assert_eq!(fold['é' as usize], bytelevel_folding::NO_FOLD); + assert_eq!(pipeline_ids(&pipeline, "xé"), vec![301, 0xA9]); + assert_eq!(pipeline_ids(&pipeline, "é"), vec![300]); + } + + /// The fold is a shortcut, never a different answer: over a real vocabulary, seeding + /// through the table and seeding plain bytes must give identical ids. Skipped when the + /// file is not fetched. + #[test] + fn fold_seeding_matches_byte_seeding_on_llama3() { + let path = "../data/llama-3-tokenizer.json"; + if !std::path::Path::new(path).exists() { + eprintln!("skip llama-3: not present (fetch with `make bench-models`)"); + return; + } + let tok = crate::Tokenizer::from_file(path).unwrap(); + let crate::models::ModelWrapper::BPE(bpe) = tok.get_model().clone() else { + panic!("llama-3 is BPE"); + }; + let folded = PipelineBPE::from_bpe(bpe.clone(), true).unwrap(); + let plain = without_fold(PipelineBPE::from_bpe(bpe, true).unwrap()); + + let Atoms::Bytes { fold, .. } = &folded.atoms else { + panic!("llama-3 is byte-level"); + }; + let entries = fold + .iter() + .filter(|&&id| id != bytelevel_folding::NO_FOLD) + .count(); + assert!(entries > 0, "no character folds in llama-3's vocabulary"); + let cjk = fold[0x3000..] + .iter() + .filter(|&&id| id != bytelevel_folding::NO_FOLD) + .count(); + eprintln!("llama-3 folds {entries} characters, {cjk} at U+3000 and above"); + + for text in [ + "朝日新聞デジタルの記事一覧です。", + "缓存的错误信息不完整，请重试。", + "Καλημέρα κόσμε, добрый день", + "mixed ascii と 日本語 and émojis 🎉🚀 nbsp\u{a0}end", + " leading spaces, ▁marks and érrors ", + ] { + assert_eq!( + pipeline_ids(&folded, text), + pipeline_ids(&plain, text), + "{text:?}" + ); + } + } + + /// Not a test: a manual probe for what the fold is worth on CJK text, model stage only, + /// cache off. Run with `cargo test --release -p tk-encode --lib fold_speed -- --ignored + /// --nocapture`. + #[test] + #[ignore] + fn fold_speed_probe_llama3_cjk() { + let path = "../data/llama-3-tokenizer.json"; + let jpn = std::fs::read_to_string("../data/fixtures/lang/jpn_Jpan.txt"); + let (Ok(tok), Ok(text)) = (crate::Tokenizer::from_file(path), jpn) else { + eprintln!("skip: llama-3 or the jpn fixture is not fetched"); + return; + }; + let crate::models::ModelWrapper::BPE(bpe) = tok.get_model().clone() else { + panic!("llama-3 is BPE"); + }; + let mut folded = PipelineBPE::from_bpe(bpe.clone(), true).unwrap(); + folded.cache_capacity = None; + let mut plain = without_fold(PipelineBPE::from_bpe(bpe, true).unwrap()); + plain.cache_capacity = None; + + // 32-character chunks as a stand-in for the pre-tokenizer's CJK runs. + let chunks: Vec<&str> = { + let mut chunks = Vec::new(); + let mut rest = text.as_str(); + while let Some((i, _)) = rest.char_indices().nth(32) { + chunks.push(&rest[..i]); + rest = &rest[i..]; + } + chunks + }; + let bytes: usize = chunks.iter().map(|c| c.len()).sum(); + + for (name, model) in [("folded", &folded), ("plain", &plain)] { + let mut scratch = model.init_scratch(); + let mut out = Vec::new(); + for _ in 0..2 { + // warm-up rep first; the second rep is the one printed + let start = std::time::Instant::now(); + for chunk in &chunks { + out.clear(); + pipeline::Model::tokenize_pipeline(model, chunk, &mut scratch, &mut out) + .unwrap(); + } + let elapsed = start.elapsed(); + eprintln!( + "{name}: {:.2} ns/B over {} chunks ({bytes} B)", + elapsed.as_nanos() as f64 / bytes as f64, + chunks.len(), + ); + } + } + } + /// A whole-word vocab hit skips the merge loop but not the vocab lookup, /// and in a byte-level vocab the words that take that path are the /// commonest ones in the text. Storing the id it found turns the next From 8da5522356743a9ee42c59f234efe13b822573f9 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:17:09 +0200 Subject: [PATCH 07/13] bench: fold_ab example for two-build process-level A/B throughput Co-Authored-By: Claude Fable 5 --- tokenizers/tk-encode/examples/fold_ab.rs | 80 ++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 tokenizers/tk-encode/examples/fold_ab.rs diff --git a/tokenizers/tk-encode/examples/fold_ab.rs b/tokenizers/tk-encode/examples/fold_ab.rs new file mode 100644 index 0000000000..65e57aa279 --- /dev/null +++ b/tokenizers/tk-encode/examples/fold_ab.rs @@ -0,0 +1,80 @@ +//! One end-to-end throughput number for one model on one corpus, for A/B runs +//! between two builds of this crate (fold on vs fold off, run alternately at +//! process level so binary-layout noise averages out). +//! +//! Mirrors `fixture_bench`'s throughput phase: single thread, ~10 kB chunks, +//! `add_special_tokens` on, one warm-up pass to fill the caches, then the +//! median of the timed passes. +//! +//! cargo run --release --example fold_ab -- [passes] + +use std::convert::TryFrom; +use std::time::Instant; + +use tk_encode::Tokenizer; +use tk_encode::pipeline::PipelineTokenizer; + +const CHUNK: usize = 10_000; +const CORPUS_CAP: usize = 4 << 20; + +fn main() { + let args: Vec = std::env::args().collect(); + let [_, model, corpus] = &args[..3] else { + eprintln!("usage: fold_ab [passes]"); + std::process::exit(2); + }; + let passes: usize = args.get(3).map_or(5, |p| p.parse().unwrap()); + + let tok = Tokenizer::from_file(model).unwrap(); + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + + let mut text = std::fs::read_to_string(corpus).unwrap(); + if text.len() > CORPUS_CAP { + let mut cap = CORPUS_CAP; + while !text.is_char_boundary(cap) { + cap -= 1; + } + text.truncate(cap); + } + let chunks: Vec<&str> = { + let mut chunks = Vec::new(); + let mut rest = text.as_str(); + while !rest.is_empty() { + let mut cut = CHUNK.min(rest.len()); + while !rest.is_char_boundary(cut) { + cut -= 1; + } + let (head, tail) = rest.split_at(cut); + chunks.push(head); + rest = tail; + } + chunks + }; + let bytes: usize = chunks.iter().map(|c| c.len()).sum(); + + let encode_all = || { + let mut ids = 0usize; + for chunk in &chunks { + ids += pipeline.encode(chunk, true).unwrap().len(); + } + ids + }; + + let ids = encode_all(); // warm-up: fills the scratch pool and the word cache + let mut timings: Vec = (0..passes) + .map(|_| { + let start = Instant::now(); + encode_all(); + start.elapsed().as_secs_f64() + }) + .collect(); + timings.sort_by(f64::total_cmp); + let median = timings[timings.len() / 2]; + + println!( + "{:.3} MB/s ({} B, {} ids, {passes} passes)", + bytes as f64 / median / 1e6, + bytes, + ids, + ); +} From 2cb2e0ebaba17016ecd6c71f66f2210e9a6a2826 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:23:47 +0200 Subject: [PATCH 08/13] perf(normalize): route every literal cut through atomsplit::Literal MetaspaceNormalizer swapped spaces by restarting a memchr search at every space and guessed the rewrite's size. It now counts the spaces first (count_matches), so the rewrite is sized exactly and a text without spaces is handed back without a scan of its own, and streams the swap through for_each_match. Normalize stage 4-7% faster on llama-2/gemma-4, e2e +3-5% (4 interleaved process rounds), ids unchanged. The drop_whitespace path sizes its rewrite from the word spans instead of the same guess. split_literal in pipeline.rs is now the one streaming literal split: Split's plain-string arm and CharDelimiterSplit both call it. CharDelimiterSplit loses its per-call span buffer (one heap Vec per chunk) and atomsplit's memchr-restart recipe struct, which the SIMD Literal scanner superseded, is deleted. Gates: 352 tk-encode + 20 atomsplit tests green, pipeline_oracle 9/9 against released 0.23.1, fmt and clippy clean (the 8 parallelism.rs dead-code warnings predate this change). Co-Authored-By: Claude Fable 5 --- tokenizers/atomsplit/src/fsm.rs | 46 ------------------- tokenizers/atomsplit/src/lib.rs | 2 +- tokenizers/atomsplit/tests/fsm.rs | 12 +---- .../tk-encode/src/normalizers/metaspace.rs | 43 +++++++++++------ .../tk-encode/src/pre_tokenizers/delimiter.rs | 21 +++++---- .../tk-encode/src/pre_tokenizers/split.rs | 20 +------- .../tk-encode/src/tokenizer/pipeline.rs | 32 ++++++++++++- 7 files changed, 77 insertions(+), 99 deletions(-) diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs index 45c8ca286b..2d7b7930a5 100644 --- a/tokenizers/atomsplit/src/fsm.rs +++ b/tokenizers/atomsplit/src/fsm.rs @@ -339,49 +339,3 @@ impl ByteLevel { fsm_byte_level(text, tags, out) } } - -/// `Split(char, Removed)` — the only pre-tokenizer that keys on a *literal char* rather than an atom -/// class, so it scans bytes directly (no classify pass). UTF-8 is self-synchronizing, so the -/// delimiter's byte pattern only matches on char boundaries. -pub struct CharDelimiterSplit(pub char); -impl CharDelimiterSplit { - /// Split on the literal char (Removed); writes spans into `out` (len ≥ `text.len()`), returns count. - #[inline] - #[must_use] - pub fn pre_tokenize(&self, text: &[u8], _tags: &mut [u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len()); - let mut buf = [0u8; 4]; - let delim = self.0.encode_utf8(&mut buf).as_bytes(); - let (n, dl) = (text.len(), delim.len()); - let (mut start, mut i, mut w) = (0usize, 0usize, 0usize); - while i + dl <= n { - // memchr the first delimiter byte, then confirm the full pattern. memchr (already a - // workspace dep) beats a scalar scan 1.4–23× here — the gap widening as the delimiter - // gets rarer over large inputs, since its SIMD skips whole 16/32/64-byte strides. - match memchr::memchr(delim[0], &text[i..n - dl + 1]) { - Some(off) if text[i + off..i + off + dl] == *delim => { - let m = i + off; - if m > start { - out[w] = Span { - start: start as u32, - end: m as u32, - }; // gap before the delimiter (Removed) - w += 1; - } - i = m + dl; - start = i; - } - Some(off) => i += off + 1, // first byte matched mid-pattern; keep scanning - None => break, - } - } - if start < n { - out[w] = Span { - start: start as u32, - end: n as u32, - }; - w += 1; - } - w - } -} diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs index 13beac00af..8882631909 100644 --- a/tokenizers/atomsplit/src/lib.rs +++ b/tokenizers/atomsplit/src/lib.rs @@ -3,7 +3,7 @@ //! One SIMD pass ([`classify`]) maps every codepoint to a tiny "atom" alphabet; a family of no-push //! FSMs ([`fsm`]) turn that atom stream into token spans (byte ranges) — the pre-tokenizer stage that //! runs before a BPE/WordPiece model. Pre-tokenizers implemented: `WhitespaceSplit`, `Punctuation`, -//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel`, `CharDelimiterSplit` (o200k and +//! `Digits`, `Whitespace`, `Bert`, `Cl100k`, `DeepSeek`, `ByteLevel` (o200k and //! Mistral's tekken are exposed as the [`fsm::fsm_o200k`] / [`fsm::fsm_tekken`] functions rather than //! recipe structs). //! diff --git a/tokenizers/atomsplit/tests/fsm.rs b/tokenizers/atomsplit/tests/fsm.rs index 1cad95e1b7..67e0c92627 100644 --- a/tokenizers/atomsplit/tests/fsm.rs +++ b/tokenizers/atomsplit/tests/fsm.rs @@ -1,8 +1,8 @@ //! Integration tests for the FSM pre-tokenizers. Kept out of `src/` so the core stays production-only. use atomsplit::classify::{classify, mask}; use atomsplit::fsm::{ - CharDelimiterSplit, Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, - fsm_deepseek, fsm_o200k, fsm_tekken, + Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, + fsm_tekken, }; /// Run a no-push fsm into a fresh buffer and return the emitted spans. @@ -57,14 +57,6 @@ fn byte_level_rules() { assert_eq!(bl("hi ok"), vec![(0, 2), (2, 4), (4, 7)]); // \s+(?!\S) leaves one space } -#[test] -fn char_delimiter_split() { - let mut out = vec![Span::default(); 8]; - // split on '/', Removed → drop delimiters, drop the empty gap between "//" - let k = CharDelimiterSplit('/').pre_tokenize(b"a/bc//d", &mut [], &mut out); - assert_eq!(&out[..k], &[(0, 1), (2, 4), (6, 7)]); -} - /// Byte-exactness gate for the class family: the NEON boundary extractor (`class_runs_into`) must equal /// the scalar run-end core (`emit_class_spans`) for every recipe, at every char-aligned truncation length so /// the < 16-byte NEON tail starts at every offset — including mid-char (chunk loop steps by 16). Corpus diff --git a/tokenizers/tk-encode/src/normalizers/metaspace.rs b/tokenizers/tk-encode/src/normalizers/metaspace.rs index b046fdb70a..193c10d835 100644 --- a/tokenizers/tk-encode/src/normalizers/metaspace.rs +++ b/tokenizers/tk-encode/src/normalizers/metaspace.rs @@ -21,6 +21,8 @@ use std::borrow::Cow; +use atomsplit::literal::Literal; + use crate::normalizers::NormalizerWrapper; use crate::normalizers::replace::{Replace, ReplacePattern}; use crate::pre_tokenizers::whitespace::WhitespaceSplit; @@ -113,12 +115,15 @@ impl pipeline::Normalizer for MetaspaceNormalizer { return Ok(Cow::Borrowed(input)); } if self.drop_whitespace { - // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space - let mut rewritten = String::with_capacity(input.len() + input.len() / 2); // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the // words back one after the other, each with its own delimiter. let mut words = Vec::new(); pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?; + // Exact when every word takes a delimiter; a word that already starts with one + // (`IfMissing`) leaves a few bytes spare. + let words_len: usize = words.iter().map(|span| span.range().len()).sum(); + let mut rewritten = + String::with_capacity(words_len + words.len() * self.delimiter.len_utf8()); for span in &words { let word = &input[span.range()]; let prepend = match self.prepend { @@ -142,23 +147,33 @@ impl pipeline::Normalizer for MetaspaceNormalizer { } PrependMode::Unconditional => true, }; + // Only spaces become delimiters; tabs and newlines are left alone. Counting them + // first sizes the rewrite exactly (a space is one byte, the delimiter up to four) + // and lets the swap stream through the batch scan instead of restarting a search + // at every space. + let space = Literal::new(b" ").expect("a space is not empty"); + let count = space.count_matches(input.as_bytes()); // Nothing to prepend and nothing to swap: hand the input back instead of copying it. - if !prepend && memchr::memchr(b' ', input.as_bytes()).is_none() { + if !prepend && count == 0 { return Ok(Cow::Borrowed(input)); } - // The delimiter is 3 bytes where a space is 1, so the rewrite grows by 2 bytes per space, hence we allocate a bit more space - let mut rewritten = String::with_capacity(input.len() + input.len() / 2); + let mut buf = [0u8; 4]; + let delimiter = self.delimiter.encode_utf8(&mut buf); + let mut rewritten = String::with_capacity( + input.len() + + (delimiter.len() - 1) * count + + if prepend { delimiter.len() } else { 0 }, + ); if prepend { - rewritten.push(self.delimiter); - } - // Only spaces become delimiters; tabs and newlines are left alone - let mut rest = input; - while let Some(space) = memchr::memchr(b' ', rest.as_bytes()) { - rewritten.push_str(&rest[..space]); - rewritten.push(self.delimiter); - rest = &rest[space + 1..]; + rewritten.push_str(delimiter); } - rewritten.push_str(rest); + let mut prev = 0; + space.for_each_match(input.as_bytes(), |start| { + rewritten.push_str(&input[prev..start]); + rewritten.push_str(delimiter); + prev = start + 1; + }); + rewritten.push_str(&input[prev..]); Ok(Cow::Owned(rewritten)) } } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs index e30f206e6b..4e2ec8f401 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/delimiter.rs @@ -1,3 +1,4 @@ +use atomsplit::literal::Literal; use serde::{Deserialize, Serialize}; use crate::pipeline; @@ -28,16 +29,18 @@ impl PreTokenizer for CharDelimiterSplit { impl pipeline::PreTokenizer for CharDelimiterSplit { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { - // native atomsplit FSM (memchr-backed single-byte scan); `Removed` — drops the delimiter, - // keeps the runs between, no empty spans. Byte-exact with the char-predicate split. - let bytes = text.as_bytes(); - let mut spans = vec![pipeline::Span::default(); bytes.len() + 1]; - let n = atomsplit::fsm::CharDelimiterSplit(self.delimiter).pre_tokenize( - bytes, - &mut [], - &mut spans, + // The same streaming literal cut as a `Split` on this char: `Removed` drops each + // delimiter and keeps the runs between them, with no empty spans. + let mut buf = [0u8; 4]; + let delimiter = Literal::new(self.delimiter.encode_utf8(&mut buf).as_bytes()) + .expect("a char is never empty"); + pipeline::split_literal( + out, + &delimiter, + text, + SplitDelimiterBehavior::Removed, + false, ); - out.extend_from_slice(&spans[..n]); Ok(()) } } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index bc5682cadd..d452a0ab1e 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -202,25 +202,9 @@ impl pipeline::PreTokenizer for Split { return Ok(()); } // A plain-string pattern streams: the batch scan hands each delimiter offset straight - // to the span fold, and nothing is built in between. The count pass sizes `out` for - // the worst case, one span per match plus the pieces around them. + // to the span fold, and nothing is built in between. if let Search::Literal(literal) = &self.search { - let width = literal.pattern().len(); - let count = literal.count_matches(text.as_bytes()); - out.reserve(2 * count + 1); - let mut fold = pipeline::SplitFold::new(out, self.behavior); - let mut prev = 0; - literal.for_each_match(text.as_bytes(), |start| { - if prev != start { - fold.segment((prev, start), self.invert); - } - fold.segment((start, start + width), !self.invert); - prev = start + width; - }); - if prev != text.len() { - fold.segment((prev, text.len()), self.invert); - } - fold.finish(); + pipeline::split_literal(out, literal, text, self.behavior, self.invert); return Ok(()); } // Not a natively-routed GPT regex either: fall back to the Regex search diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 22fbc84c89..097af92ae8 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1030,9 +1030,39 @@ pub fn split_matches( fold.finish(); } +/// [`split_matches`] for a plain-byte pattern, shared by every pre-tokenizer that cuts on a +/// literal (`Split` with a string pattern, `CharDelimiterSplit`). The count pass sizes `out` +/// for the worst case, one span per match plus the pieces around them; the batch scan then +/// hands each match offset straight to the fold, so no segment vector is built. `invert` +/// flips which side counts as the delimiter, as in `Split`. +pub(crate) fn split_literal( + out: &mut Vec, + literal: &atomsplit::literal::Literal, + text: &str, + behavior: SplitDelimiterBehavior, + invert: bool, +) { + let width = literal.pattern().len(); + let count = literal.count_matches(text.as_bytes()); + out.reserve(2 * count + 1); + let mut fold = SplitFold::new(out, behavior); + let mut prev = 0; + literal.for_each_match(text.as_bytes(), |start| { + if prev != start { + fold.segment((prev, start), invert); + } + fold.segment((start, start + width), !invert); + prev = start + width; + }); + if prev != text.len() { + fold.segment((prev, text.len()), invert); + } + fold.finish(); +} + /// The streaming form of [`split_matches`]: feed it the covering `(offsets, is_match)` /// segments left to right with [`SplitFold::segment`], then call [`SplitFold::finish`]. -/// Callers that produce segments one at a time (the literal path in `Split`) drive it +/// Callers that produce segments one at a time ([`split_literal`]) drive it /// directly and never build the segment vector. /// /// One span under construction and the previous segment's flag are the only state. A span From 809246208f117254b7920fc969256df309efc714 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:41:08 +0200 Subject: [PATCH 09/13] perf(cache): pack the word key with overlapping reads, not a buffer copy pack_word built its u128 key by zeroing a 16-byte buffer and copying the word in with a variable length, which LLVM lowers to a memset and a memcpy libcall. One pair per looked-up word put a third of gpt2's whole encode inside libsystem_platform (samply: 27.5% memset + 6.4% memmove, both under WordCache::lookup). Two fixed-width reads, one from each end of the word, produce the same value with no calls: the overlapped middle bytes are identical, so or-ing the halves is harmless, and every byte above the length stays zero for key equality. pack_word_matches_the_buffer_form pins the new form to the old one for every packable length and byte pattern. Measured (interleaved process rounds, ids identical, oracle 9/9): gpt2 +38/+30% (eng/agentic), llama-3 +39/+33%, glm +33/+31%, gpt-oss +23/+19%, mistral +5/+21%, llama-2 +16/+5%, gemma-4 +15/+12%. An 8-way co-runner test holds per-process throughput within 3% of solo, so the encode loop stays core-bound, not memory-bound. Co-Authored-By: Claude Fable 5 --- tokenizers/tk-encode/src/utils/word_cache.rs | 44 ++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 96add857dc..def3c129ad 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -642,9 +642,28 @@ fn pack_word(word: &[u8]) -> Option { if len == 0 || len > 15 { return None; } - let mut lanes = [0u8; 16]; - lanes[..len].copy_from_slice(word); - Some(u128::from_le_bytes(lanes) | ((len as u128) << 120)) + // Overlapping fixed-width reads instead of a zeroed buffer and a + // variable-length copy: the buffer form compiled to a memset and a memcpy + // libcall per word, which the profile put at a third of gpt2's encode. Two + // reads of the same width, one from each end of the word, cover every + // length; the overlapped middle bytes are identical, so or-ing them is + // harmless, and every byte above `len` stays zero for key equality. + let value = if len >= 8 { + let head = u64::from_le_bytes(word[..8].try_into().unwrap()) as u128; + let tail = u64::from_le_bytes(word[len - 8..].try_into().unwrap()) as u128; + head | (tail << ((len - 8) * 8)) + } else if len >= 4 { + let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u128; + let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u128; + head | (tail << ((len - 4) * 8)) + } else { + // 1 to 3 bytes: first, middle and last byte land on their own lanes + // (for the shorter lengths some of the three are the same byte). + (word[0] as u128) + | ((word[len / 2] as u128) << (len / 2 * 8)) + | ((word[len - 1] as u128) << ((len - 1) * 8)) + }; + Some(value | ((len as u128) << 120)) } // ---------------------------------------------------------------- overflow storage @@ -764,6 +783,25 @@ mod tests { } } + /// Every packable length and byte pattern packs as if the word were copied + /// into a zeroed 16-byte buffer: same bytes in the low lanes, zeros above + /// them, the length in the top lane. Pins [`pack_word`]'s read trickery to + /// the plain buffer form it replaced. + #[test] + fn pack_word_matches_the_buffer_form() { + for len in 0..=16usize { + for fill in [0x00u8, 0x5a, 0xff] { + let word: Vec = (0..len).map(|i| fill.wrapping_add(i as u8)).collect(); + let expected = (len >= 1 && len <= 15).then(|| { + let mut lanes = [0u8; 16]; + lanes[..len].copy_from_slice(&word); + u128::from_le_bytes(lanes) | ((len as u128) << 120) + }); + assert_eq!(pack_word(&word), expected, "len {len} fill {fill:#x}"); + } + } + } + #[test] fn roundtrip() { let mut cache = WordCache::new(1 << 8); From 459be6c7fa57da7ec2cc3d98c06e73302e2656ec Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 02:15:29 +0200 Subject: [PATCH 10/13] perf(pipeline): fuse the native FSM scans into the model step The staged path collects every span of a chunk into a buffer, then walks the buffer calling the model through the per-span (model, scratch) dispatch match: a span-buffer round trip and a dispatch per word. Every FSM-routed pre-tokenizer over a BPE model (the gpt2, cl100k-family, o200k, tekken and deepseek shapes: all eight byte-level benchmark models but the two SentencePiece ones) now runs fused instead: encode drives the scan_* form of the FSM, which hands each span to the model the moment it is cut. No span buffer, dispatch settled once per chunk, and the model reads each word while its bytes are still hot from the scan. Every fsm_* keeps its no-push slice form as a thin wrapper over the new scan_* (one emit closure writing spans in place), so the existing byte-exactness tests keep pinning both forms at once. PipelineBPE::tokenize_span is the per-word body factored out of the Model trait impl, so the fused and staged paths share one body and the ignore_merges branch keeps living in exactly one place. Routing asks the pre-tokenizer itself: Split::native_fsm and Sequence::native_fsm apply the same recognition their pre_tokenize applies (deepseek's three-Split composition, the identity-children collapse), so the fused check cannot disagree with the split it bypasses. Measured with a flat encode loop, interleaved process rounds (eng_Latn / agentic_swe): gpt2 +17%, llama-3 +15/+16%, glm-5.2 +14/+15%, gpt-oss +14/+15%, mistral-small-4 +12/+17%, deepseek-v4 +11/+9%. llama-2 and gemma-4 do not route (ProvenCuts / literal splits) and are unchanged. Ids identical on every fixture; oracle 9/9. Co-Authored-By: Claude Fable 5 --- tokenizers/atomsplit/src/fsm.rs | 8 +- tokenizers/atomsplit/src/fsm/byte_level.rs | 27 +++-- tokenizers/atomsplit/src/fsm/cl100k.rs | 27 +++-- tokenizers/atomsplit/src/fsm/deepseek.rs | 73 ++++++----- tokenizers/atomsplit/src/fsm/o200k.rs | 83 +++++++------ tokenizers/tk-encode/src/models/bpe/model.rs | 31 +++-- .../tk-encode/src/pre_tokenizers/sequence.rs | 20 ++++ .../tk-encode/src/pre_tokenizers/split.rs | 8 ++ .../tk-encode/src/tokenizer/pipeline.rs | 113 ++++++++++++++++-- tokenizers/tk-encode/src/utils/word_cache.rs | 2 +- 10 files changed, 272 insertions(+), 120 deletions(-) diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs index 2d7b7930a5..fe85b04005 100644 --- a/tokenizers/atomsplit/src/fsm.rs +++ b/tokenizers/atomsplit/src/fsm.rs @@ -238,10 +238,10 @@ mod byte_level; mod cl100k; mod deepseek; mod o200k; -pub use byte_level::fsm_byte_level; -pub use cl100k::{fsm_cl100k, fsm_cl100k_cap}; -pub use deepseek::fsm_deepseek; -pub use o200k::{fsm_o200k, fsm_tekken}; +pub use byte_level::{fsm_byte_level, scan_byte_level}; +pub use cl100k::{fsm_cl100k, fsm_cl100k_cap, scan_cl100k_cap}; +pub use deepseek::{fsm_deepseek, scan_deepseek}; +pub use o200k::{fsm_o200k, fsm_tekken, scan_o200k, scan_tekken}; // ── Composition recipes ──────────────────────────────────────────────────────────────────────── // Each pre-tokenizer = (classify → fsm shape + params). `tags` and `out` are caller-owned diff --git a/tokenizers/atomsplit/src/fsm/byte_level.rs b/tokenizers/atomsplit/src/fsm/byte_level.rs index 73f84ac9fe..93f3776312 100644 --- a/tokenizers/atomsplit/src/fsm/byte_level.rs +++ b/tokenizers/atomsplit/src/fsm/byte_level.rs @@ -8,6 +8,19 @@ use super::*; #[must_use] pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_byte_level(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The scan under [`fsm_byte_level`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_byte_level(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { + debug_assert!(tags.len() >= text.len()); let end = text.len(); // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) @@ -27,7 +40,6 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { }; let mut i = 0; - let mut w = 0usize; while i < end { let start = i; match tags[i] & 0x0F { @@ -68,14 +80,9 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. _ => i += char_len(text[i]), } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; + emit(Span { + start: start as u32, + end: i as u32, + }); } - w } diff --git a/tokenizers/atomsplit/src/fsm/cl100k.rs b/tokenizers/atomsplit/src/fsm/cl100k.rs index 65976ecef8..95f65aca14 100644 --- a/tokenizers/atomsplit/src/fsm/cl100k.rs +++ b/tokenizers/atomsplit/src/fsm/cl100k.rs @@ -14,10 +14,19 @@ pub fn fsm_cl100k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { #[must_use] pub fn fsm_cl100k_cap(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); - cl100k(text, tags, out, digit_cap) + let mut w = 0usize; + scan_cl100k_cap(text, tags, digit_cap, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w } -fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize { +/// The scan under [`fsm_cl100k_cap`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_cl100k_cap(text: &[u8], tags: &[u8], digit_cap: usize, mut emit: impl FnMut(Span)) { + debug_assert!(tags.len() >= text.len()); // Leading-atom values, as `const` so the `match` below is a dense jump table (not an if-cascade): // the dispatch is O(1) and a token never pays for a rule it can't start (e.g. non-number tokens // never test the number rule — which is what the POC's const-gating removed by hand; here it's free). @@ -41,7 +50,6 @@ fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; let mut i = 0; - let mut w = 0usize; while i < end { let start = i; let b = text[i]; @@ -105,14 +113,9 @@ fn cl100k(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usize) -> usize // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. _ => i += char_len(b), } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; + emit(Span { + start: start as u32, + end: i as u32, + }); } - w } diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs index f54afb8fac..75d3131e61 100644 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ b/tokenizers/atomsplit/src/fsm/deepseek.rs @@ -39,7 +39,20 @@ fn ds_is_cjk_at(text: &[u8], p: usize) -> bool { /// `\w` but categorically `\p{S}`) take the `[\p{P}\p{S}]` path, not the letter run. #[must_use] pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + debug_assert!(out.len() >= text.len()); + let mut w = 0usize; + scan_deepseek(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The scan under [`fsm_deepseek`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_deepseek(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { + debug_assert!(tags.len() >= text.len()); // Leading-atom values as `const` → the `match` is a dense jump table (see `cl100k`). The Split // precedence (digits → CJK → big-regex alts) is preserved because the atom partition is disjoint. // `Mark` refined as an Other_Alphabetic symbol (Ⓘ …): coarse `LETTER_MARK`, but categorically `\p{S}` @@ -113,7 +126,6 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { }; let mut i = 0; - let mut w = 0usize; while i < end { let start = i; let b = text[i]; @@ -129,13 +141,10 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { { p += 3; } - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; - w += 1; + emit(Span { + start: start as u32, + end: p as u32, + }); i = p; continue; } @@ -150,31 +159,22 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { } if is_lm(p) { if last > i { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: last as u32, - } - }; // gap sans prefix char - w += 1; + emit(Span { + start: start as u32, + end: last as u32, + }); // gap sans prefix char } let e = letter_run(p); - unsafe { - *out.get_unchecked_mut(w) = Span { - start: last as u32, - end: e as u32, - } - }; // prefix char + `[\p{L}\p{M}]+` - w += 1; + emit(Span { + start: last as u32, + end: e as u32, + }); // prefix char + `[\p{L}\p{M}]+` i = e; } else { - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: p as u32, - } - }; // whole gap run is one piece - w += 1; + emit(Span { + start: start as u32, + end: p as u32, + }); // whole gap run is one piece i = p; } continue; @@ -234,14 +234,9 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. _ => i += char_len(b), } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; + emit(Span { + start: start as u32, + end: i as u32, + }); } - w } diff --git a/tokenizers/atomsplit/src/fsm/o200k.rs b/tokenizers/atomsplit/src/fsm/o200k.rs index 28e15640a3..de02fbd79f 100644 --- a/tokenizers/atomsplit/src/fsm/o200k.rs +++ b/tokenizers/atomsplit/src/fsm/o200k.rs @@ -50,19 +50,18 @@ fn o200k_letter_match(tags: &[u8], p: usize, re: usize) -> usize { e } -/// Emit the o200k case-split of the letter run `[ls, re)` into `out[*w..]`: the first sub-token starts at +/// Emit the o200k case-split of the letter run `[ls, re)`: the first sub-token starts at /// `pfx` (the optional `[^\r\n\p{L}\p{N}]?` prefix; `pfx == ls` when none), the last absorbs a trailing /// contraction (`CONTRACTION` — off for tekken). Returns the new cursor (past the contraction). /// `ls < re` (caller-guaranteed). #[inline(always)] -fn emit_o200k_letters( +fn emit_o200k_letters( text: &[u8], tags: &[u8], pfx: usize, ls: usize, re: usize, - out: &mut [Span], - w: &mut usize, + emit: &mut E, ) -> usize { let (mut p, mut first, mut cursor) = (ls, true, re); while p < re { @@ -73,13 +72,10 @@ fn emit_o200k_letters( } else { e }; - unsafe { - *out.get_unchecked_mut(*w) = Span { - start: start as u32, - end: tok_end as u32, - } - }; - *w += 1; + emit(Span { + start: start as u32, + end: tok_end as u32, + }); first = false; cursor = tok_end; p = e; @@ -94,22 +90,47 @@ fn emit_o200k_letters( /// Unlike deepseek there are no gaps: rule 4's `[^\s\p{L}\p{N}]+` is a catch-all. Scalar; ┌ OWNER: shared ┐ #[must_use] pub fn fsm_o200k(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) + debug_assert!(out.len() >= text.len()); + let mut w = 0usize; + scan_o200k(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w } /// Mistral tekken ([`crate::regexes::TEKKEN`]) — the o200k FSM with the contraction suffix off and one /// token per digit. Every other rule is shared, so both are the same code monomorphized twice. #[must_use] pub fn fsm_tekken(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { - o200k::(text, tags, out) + debug_assert!(out.len() >= text.len()); + let mut w = 0usize; + scan_tekken(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The scan under [`fsm_o200k`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_o200k(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + o200k::(text, tags, emit); +} + +/// The scan under [`fsm_tekken`]; see [`scan_o200k`]. +pub fn scan_tekken(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + o200k::(text, tags, emit); } -fn o200k( +fn o200k( text: &[u8], tags: &[u8], - out: &mut [Span], -) -> usize { - debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + mut emit: E, +) { + debug_assert!(tags.len() >= text.len()); let end = text.len(); // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) @@ -168,12 +189,11 @@ fn o200k( // rules 5-7 (`\s*[\r\n]+ | \s+(?!\S) | \s+`) → the shared `ws_tail` (identical to cl100k). let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; // The letter rules: case-split the run starting at `ls`, first sub-token starting at the prefix `pfx`. - let letters = |pfx: usize, ls: usize, out: &mut [Span], w: &mut usize| -> usize { - emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), out, w) + let letters = |pfx: usize, ls: usize, emit: &mut E| -> usize { + emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), emit) }; let mut i = 0; - let mut w = 0usize; while i < end { let start = i; let b = text[i]; @@ -191,12 +211,12 @@ fn o200k( // take the `[^\r\n\p{L}\p{N}]?` prefix / rule-4 path instead. LET | MRK => { if tags[i] != ASM && tags[i] != ZWJ { - i = letters(i, i, out, &mut w); + i = letters(i, i, &mut emit); continue; } let a = i + char_len(b); if is_lm(a) { - i = letters(i, a, out, &mut w); + i = letters(i, a, &mut emit); continue; } i = other(i); // ∈ NOT_WS_L_N ⇒ > i @@ -205,7 +225,7 @@ fn o200k( SPC => { let a = i + 1; // Space is ASCII (0x20) if is_lm(a) { - i = letters(i, a, out, &mut w); + i = letters(i, a, &mut emit); continue; } let p = other(a); @@ -215,7 +235,7 @@ fn o200k( WSO => { let a = i + char_len(b); if is_lm(a) { - i = letters(i, a, out, &mut w); + i = letters(i, a, &mut emit); continue; } i = ws(i); @@ -225,7 +245,7 @@ fn o200k( CON | PUN | APO | SYM | NMO | CTL => { let a = i + char_len(b); if is_lm(a) { - i = letters(i, a, out, &mut w); + i = letters(i, a, &mut emit); continue; } i = other(i); // ∈ NOT_WS_L_N ⇒ > i @@ -233,14 +253,9 @@ fn o200k( // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. _ => i += char_len(b), } - // SAFETY: tokens partition the input, so `w < #tokens <= end < out.len()` (out ≥ text.len()+? ; callers size n+1). - unsafe { - *out.get_unchecked_mut(w) = Span { - start: start as u32, - end: i as u32, - } - }; - w += 1; + emit(Span { + start: start as u32, + end: i as u32, + }); } - w } diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 1ab2d2ae78..8fdec30d5e 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -856,17 +856,20 @@ impl PipelineBPE { } } -impl pipeline::Model for PipelineBPE { - type Scratch = BpeScratch; - - fn tokenize_pipeline( +impl PipelineBPE { + /// One pre-token through the cache and the merge loop. This is the whole model + /// step for a span and it cannot fail; it stands alone so the fused byte-level + /// scan can call it per span as the span is cut, and + /// [`Model::tokenize_pipeline`](pipeline::Model::tokenize_pipeline) is the same + /// body behind the trait. + pub(crate) fn tokenize_span( &self, sequence: &str, - scratch: &mut Self::Scratch, + scratch: &mut BpeScratch, output: &mut Vec, - ) -> Result<()> { + ) { if sequence.is_empty() { - return Ok(()); + return; } let BpeScratch { @@ -881,7 +884,7 @@ impl pipeline::Model for PipelineBPE { match cache.lookup(sequence.as_bytes()) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); - return Ok(()); + return; } Lookup::Miss(at) => placement = at, } @@ -902,7 +905,19 @@ impl pipeline::Model for PipelineBPE { { cache.insert(at, output[start..].iter().map(|token| token.id)); } + } +} +impl pipeline::Model for PipelineBPE { + type Scratch = BpeScratch; + + fn tokenize_pipeline( + &self, + sequence: &str, + scratch: &mut Self::Scratch, + output: &mut Vec, + ) -> Result<()> { + self.tokenize_span(sequence, scratch, output); Ok(()) } diff --git a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs index b01d4686d8..d216843e27 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/sequence.rs @@ -62,6 +62,26 @@ impl PipelineSequence { /// Isolated, non-inverted `Split`s carrying deepseek's `[\p{N}{1,3}, CJK, big]` regexes (the trailing /// byte-map `ByteLevel` converts to `PipelinePreTokenizer::None`). Routes the whole split to one /// `fsm_deepseek` pass. + /// The single native FSM this sequence reduces to, if any: deepseek's three-Split + /// composition whole, or the lone real child left once identity passes + /// (`PipelinePreTokenizer::None`) are skipped. Both answers mirror the routing + /// `pre_tokenize` below applies, in the same order. + pub(crate) fn native_fsm(&self) -> Option { + if self.is_deepseek() { + return Some(pipeline::FusedScan::DeepSeek); + } + let mut work = self + .pre_tokenizers + .iter() + .filter(|c| !matches!(c, PipelinePreTokenizer::None)); + match (work.next(), work.next()) { + (Some(PipelinePreTokenizer::Split(split)), None) => { + split.native_fsm().map(pipeline::FusedScan::Gpt) + } + _ => None, + } + } + fn is_deepseek(&self) -> bool { use crate::pre_tokenizers::split::SplitPattern; use crate::tokenizer::SplitDelimiterBehavior::Isolated; diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index d452a0ab1e..a38a7b6216 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -126,6 +126,14 @@ impl Split { }) } + /// The native atomsplit FSM this split routes to, if its shape allows it: the + /// same recognition and `Isolated`, not-inverted filter `pre_tokenize` applies, + /// so a caller planning around the FSM cannot disagree with the split itself. + pub(crate) fn native_fsm(&self) -> Option { + self.fsm + .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated) + } + /// Pipeline canonicalization. A recognized whole-covering GPT regex shipped /// as `(invert=true, behavior=Removed)` — the tiktoken-conversion convention /// used by cl100k/o200k — is byte-exactly equivalent to `(invert=false, diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 097af92ae8..ec2d397e3c 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -12,6 +12,7 @@ use crate::models::wordlevel::WordLevel; use crate::models::wordpiece::{PipelineWordPiece, WordPieceScratch}; use crate::processors::bert::BertProcessing; use crate::processors::roberta::RobertaProcessing; +use crate::utils::GptFsm; use crate::utils::byte_level::GPT2_REGEX_STR; use crate::vocab::bucket_added_vocabulary::{ AddedToken as BucketAddedToken, AddedVocabulary as BucketAddedVocabulary, @@ -150,6 +151,28 @@ pub enum PipelinePreTokenizer { None, } +/// A pre-tokenizer whose whole split runs as one native FSM pass, named ahead of time +/// so [`PipelineTokenizer::encode_fused`] can drive the scan and the model together. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(crate) enum FusedScan { + Gpt(GptFsm), + DeepSeek, +} + +impl PipelinePreTokenizer { + /// The single native FSM this whole pre-tokenizer reduces to, if any: a recognized + /// lone `Split`, or a `Sequence` that collapses to one. Each arm asks the child for + /// the routing its own `pre_tokenize` applies, so a fused caller cannot disagree + /// with the split it bypasses. + pub(crate) fn native_fsm(&self) -> Option { + match self { + Self::Split(split) => split.native_fsm().map(FusedScan::Gpt), + Self::Sequence(sequence) => sequence.native_fsm(), + _ => None, + } + } +} + impl PreTokenizer for PipelinePreTokenizer { fn pre_tokenize(&self, text: &str, out: &mut Vec) -> Result<()> { match self { @@ -816,18 +839,24 @@ impl PipelineTokenizer { } Segment::Text(normalized_chunk) => { if STAGE >= Self::STAGE_SPLIT { - // Pre-tokenize the chunk of normalized text - pre_tokens.clear(); - self.pre_tokenizer - .pre_tokenize(normalized_chunk, pre_tokens)?; - if STAGE >= Self::STAGE_MODEL { - // Tokenize each chunk - for pre_token in pre_tokens.iter() { - self.model.tokenize_pipeline( - &normalized_chunk[pre_token.range()], - scratch, - output, - )?; + if STAGE >= Self::STAGE_MODEL + && self.encode_fused(normalized_chunk, scratch, output) + { + // Split and model ran as one pass; nothing left to do. + } else { + // Pre-tokenize the chunk of normalized text + pre_tokens.clear(); + self.pre_tokenizer + .pre_tokenize(normalized_chunk, pre_tokens)?; + if STAGE >= Self::STAGE_MODEL { + // Tokenize each chunk + for pre_token in pre_tokens.iter() { + self.model.tokenize_pipeline( + &normalized_chunk[pre_token.range()], + scratch, + output, + )?; + } } } } @@ -843,6 +872,66 @@ impl PipelineTokenizer { } Ok(()) } + + /// The fused fast path: cut `chunk` with its native FSM and hand each span + /// straight to the BPE model the moment it is emitted. The span buffer, the + /// per-span model dispatch and the per-span trait call of the staged path all + /// disappear, and the model reads each word while its bytes are still hot from + /// the scan. Returns `false` when this tokenizer is not that shape, and the + /// staged path must run instead. + fn encode_fused( + &self, + chunk: &str, + scratch: &mut PipelineModelScratch, + output: &mut Vec, + ) -> bool { + let Some(scan) = self.pre_tokenizer.native_fsm() else { + return false; + }; + let PipelineModel::BPE(model) = &self.model else { + return false; + }; + let PipelineModelScratch::BPE(scratch) = scratch else { + return false; + }; + thread_local! { + static TAGS: RefCell> = const { RefCell::new(Vec::new()) }; + } + let bytes = chunk.as_bytes(); + TAGS.with(|cell| { + let tags = &mut *cell.borrow_mut(); + if tags.len() < bytes.len() { + tags.resize(bytes.len(), 0); // grow-only, as in `classify_into_spans` + } + classify(bytes, &mut tags[..bytes.len()]); + let tags = &tags[..bytes.len()]; + use atomsplit::fsm::{ + scan_byte_level, scan_cl100k_cap, scan_deepseek, scan_o200k, scan_tekken, + }; + // One emit closure literal per arm: each scan gets its own instance by + // value, which is what lets the emit inline into the scan loop. + match scan { + FusedScan::Gpt(GptFsm::Gpt2) => scan_byte_level(bytes, tags, |span| { + model.tokenize_span(&chunk[span.range()], scratch, output); + }), + FusedScan::Gpt(GptFsm::Cl100k { digit_cap }) => { + scan_cl100k_cap(bytes, tags, digit_cap, |span| { + model.tokenize_span(&chunk[span.range()], scratch, output); + }) + } + FusedScan::Gpt(GptFsm::O200k) => scan_o200k(bytes, tags, |span| { + model.tokenize_span(&chunk[span.range()], scratch, output); + }), + FusedScan::Gpt(GptFsm::Tekken) => scan_tekken(bytes, tags, |span| { + model.tokenize_span(&chunk[span.range()], scratch, output); + }), + FusedScan::DeepSeek => scan_deepseek(bytes, tags, |span| { + model.tokenize_span(&chunk[span.range()], scratch, output); + }), + } + }); + true + } } /// Streaming decoder over a [`PipelineTokenizer`]; see [`PipelineTokenizer::decode_stream`]. diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index def3c129ad..bc44147819 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -792,7 +792,7 @@ mod tests { for len in 0..=16usize { for fill in [0x00u8, 0x5a, 0xff] { let word: Vec = (0..len).map(|i| fill.wrapping_add(i as u8)).collect(); - let expected = (len >= 1 && len <= 15).then(|| { + let expected = (1..=15).contains(&len).then(|| { let mut lanes = [0u8; 16]; lanes[..len].copy_from_slice(&word); u128::from_le_bytes(lanes) | ((len as u128) << 120) From 24813cf6dec8ee6ffd5fdd2557b5118a94424dcb Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 12:18:42 +0200 Subject: [PATCH 11/13] perf(pretok): boundary-mask scanners for the regex-shaped FSMs Replace the per-run scalar walk of every regex-shaped pre-tokenizer FSM (byte-level, cl100k at digit caps 1/3/unbounded, o200k, tekken, deepseek) with 64-byte batch scanners behind one shared walker: per- class bitmasks built from the classify tag stream, token starts computed by shifted-mask algebra in u64 registers (transcribed per scheme from gigatoken's scanners, MIT), and everything the algebra cannot prove locally deferred to the scheme's scalar advance, extracted from each scan so both paths share one body. Feeding the masks from tags instead of raw bytes keeps non-ASCII batches on the fast path. Bad zones: char-counted digit groups over multi-byte chars, mark chars (run- contextual class), upper-after-caseless case splits, contraction and whitespace batch-edge straddles, multi-byte whitespace, and any deepseek batch containing a CJK-range char (tested before the tag masks are built, so CJK text degenerates to the scalar scan plus one test). The walker interleaves proven starts with bad zones: a span must never be emitted across one. The per-target surface is one 64-byte block classifier each for NEON, SSE2 and wasm32 simd128; other targets fall back to the scalar scans. FSM stage on eng_Latn/agentic_swe: byte_level 3.5x/2.5x, cl100k 2.8x/1.9x, o200k 3.2x/1.9x, tekken 3.3x/2.1x, deepseek 3.1x/1.7x; cmn_Hani 1.2-1.9x with deepseek at parity by construction. Encode e2e, process-alternated rounds with identical token streams: gpt2 249 -> 321 MB/s eng (+29%), +11% swe; llama-3 +28%/+16% (cmn -1.6%, noise); glm-5.2 +27%/+13%; gpt-oss +49%/+26%; mistral +46%/+25%; deepseek +52%/+23% (fancy-regex builds both sides). Byte-exactness pinned per scheme by padding-sweep differentials (every 64-byte-edge offset, also run on x86_64) and byte-exact spans over three full fixtures; pipeline_oracle 9/9 vs the released crate. Co-Authored-By: Claude Fable 5 --- tokenizers/atomsplit/src/fsm.rs | 10 +- tokenizers/atomsplit/src/fsm/byte_level.rs | 102 ++-- tokenizers/atomsplit/src/fsm/cl100k.rs | 149 +++--- tokenizers/atomsplit/src/fsm/deepseek.rs | 220 +++++--- tokenizers/atomsplit/src/fsm/masked.rs | 440 ++++++++++++++++ tokenizers/atomsplit/src/fsm/masked/block.rs | 427 ++++++++++++++++ .../atomsplit/src/fsm/masked/byte_level.rs | 146 ++++++ tokenizers/atomsplit/src/fsm/masked/cl100k.rs | 313 ++++++++++++ .../atomsplit/src/fsm/masked/deepseek.rs | 306 +++++++++++ tokenizers/atomsplit/src/fsm/masked/o200k.rs | 473 ++++++++++++++++++ tokenizers/atomsplit/src/fsm/o200k.rs | 191 +++++-- tokenizers/atomsplit/tests/fsm.rs | 145 +++++- .../tk-encode/src/tokenizer/pipeline.rs | 13 +- 13 files changed, 2694 insertions(+), 241 deletions(-) create mode 100644 tokenizers/atomsplit/src/fsm/masked.rs create mode 100644 tokenizers/atomsplit/src/fsm/masked/block.rs create mode 100644 tokenizers/atomsplit/src/fsm/masked/byte_level.rs create mode 100644 tokenizers/atomsplit/src/fsm/masked/cl100k.rs create mode 100644 tokenizers/atomsplit/src/fsm/masked/deepseek.rs create mode 100644 tokenizers/atomsplit/src/fsm/masked/o200k.rs diff --git a/tokenizers/atomsplit/src/fsm.rs b/tokenizers/atomsplit/src/fsm.rs index fe85b04005..5d6b950d4e 100644 --- a/tokenizers/atomsplit/src/fsm.rs +++ b/tokenizers/atomsplit/src/fsm.rs @@ -6,8 +6,8 @@ //! [`class_runs_into`]: on aarch64/wasm the SIMD movemask boundary-extractor + homogeneous-chunk //! early-out (in `simd_fsm`), elsewhere the scalar run-end core ([`emit_class_spans`]). The //! regex-shaped ones ([`fsm_cl100k`] / [`fsm_o200k`] / [`fsm_tekken`] / [`fsm_deepseek`] / -//! [`fsm_byte_level`]) are scalar jump-tables (only the class family's [`class_runs_into`] has a SIMD -//! path). +//! [`fsm_byte_level`]) are scalar jump-tables; byte_level additionally has a boundary-mask SIMD +//! form ([`scan_byte_level_masked`], aarch64 only). pub(crate) use crate::classify::{Atom, char_len, classify, in_mask, mask}; // Atom-tag aliases, shared with the per-tokenizer FSM submodules (`fsm/*.rs`) via `use super::*`. @@ -237,10 +237,16 @@ pub fn emit_class_spans( mod byte_level; mod cl100k; mod deepseek; +mod masked; mod o200k; pub use byte_level::{fsm_byte_level, scan_byte_level}; pub use cl100k::{fsm_cl100k, fsm_cl100k_cap, scan_cl100k_cap}; pub use deepseek::{fsm_deepseek, scan_deepseek}; +pub use masked::{ + fsm_byte_level_masked, fsm_cl100k_cap_masked, fsm_deepseek_masked, fsm_o200k_masked, + fsm_tekken_masked, scan_byte_level_masked, scan_cl100k_cap_masked, scan_deepseek_masked, + scan_o200k_masked, scan_tekken_masked, +}; pub use o200k::{fsm_o200k, fsm_tekken, scan_o200k, scan_tekken}; // ── Composition recipes ──────────────────────────────────────────────────────────────────────── diff --git a/tokenizers/atomsplit/src/fsm/byte_level.rs b/tokenizers/atomsplit/src/fsm/byte_level.rs index 93f3776312..479a4b1de3 100644 --- a/tokenizers/atomsplit/src/fsm/byte_level.rs +++ b/tokenizers/atomsplit/src/fsm/byte_level.rs @@ -17,14 +17,12 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { w } -/// The scan under [`fsm_byte_level`]: hands each token to `emit` the moment it is cut, -/// so a caller can consume tokens in place instead of collecting a span buffer first. -pub fn scan_byte_level(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { - debug_assert!(tags.len() >= text.len()); - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; +/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of the +/// byte-level regex. [`scan_byte_level`] loops it over the whole text; the masked scanner +/// ([`super::scan_byte_level_masked`]) re-derives tokens with it where its batch masks are not +/// trustworthy. +#[inline(always)] +pub(super) fn advance_byte_level(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { // `\s+(?!\S)|\s+`: the whole run at EOF, else leave the last ws char for the next ` ?`-prefixed run. let ws = |i: usize| -> usize { let re = run_end(tags, i, end, mask::WS); @@ -38,51 +36,57 @@ pub fn scan_byte_level(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { if last > i { last } else { re } } }; - - let mut i = 0; - while i < end { - let start = i; - match tags[i] & 0x0F { - LET => i = run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm) - NW | NO => i = run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED - MRK | CON | PUN | SYM | NMO | CTL => i = run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+` - // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set) - APO => { - let adv = match (text.get(i + 1), text.get(i + 2)) { - (Some(b's' | b't' | b'm' | b'd'), _) => 2, - (Some(b'r'), Some(b'e')) - | (Some(b'v'), Some(b'e')) - | (Some(b'l'), Some(b'l')) => 3, - _ => 0, - }; - i = if adv > 0 { - i + adv - } else { - run_end(tags, i, end, mask::NOT_WS_L_N) - }; + match tags[i] & 0x0F { + LET => run_end(tags, i, end, mask::LETTER), // ` ?\p{L}+` (space taken by the Space arm) + NW | NO => run_end(tags, i, end, mask::NUMBER), // ` ?\p{N}+` — UNBOUNDED + MRK | CON | PUN | SYM | NMO | CTL => run_end(tags, i, end, mask::NOT_WS_L_N), // ` ?[^…]+` + // `'s|'t|'re|'ve|'m|'ll|'d` (case-sensitive), else `[^\s\p{L}\p{N}]+` (apostrophe ∈ that set) + APO => { + let adv = match (text.get(i + 1), text.get(i + 2)) { + (Some(b's' | b't' | b'm' | b'd'), _) => 2, + (Some(b'r'), Some(b'e')) | (Some(b'v'), Some(b'e')) | (Some(b'l'), Some(b'l')) => 3, + _ => 0, + }; + if adv > 0 { + i + adv + } else { + run_end(tags, i, end, mask::NOT_WS_L_N) } - // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run, - // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run). - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = match tags.get(a).map(|&t| t & 0x0F) { - Some(LET) => run_end(tags, a, end, mask::LETTER), - Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER), - Some(t) if in_mask(t, mask::NOT_WS_L_N) => { - run_end(tags, a, end, mask::NOT_WS_L_N) - } - _ => ws(i), - }; + } + // Space: the ` ?` prefix — attach one space to a following letter / number / "other" run, + // else it's whitespace (rules `\s+(?!\S)|\s+`, which leave one space for the next run). + SPC => { + let a = i + 1; // Space is ASCII (0x20) + match tags.get(a).map(|&t| t & 0x0F) { + Some(LET) => run_end(tags, a, end, mask::LETTER), + Some(NW) | Some(NO) => run_end(tags, a, end, mask::NUMBER), + Some(t) if in_mask(t, mask::NOT_WS_L_N) => run_end(tags, a, end, mask::NOT_WS_L_N), + _ => ws(i), } - // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines - // never prefix a run. - WSO | NLN => i = ws(i), - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(text[i]), } + // WsOther / Newline: whitespace only — the ` ?` prefix is a literal 0x20, so tabs/newlines + // never prefix a run. + WSO | NLN => ws(i), + // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. + _ => i + char_len(text[i]), + } +} + +/// The scan under [`fsm_byte_level`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_byte_level(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { + debug_assert!(tags.len() >= text.len()); + let end = text.len(); + // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior + // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) + let tags = &tags[..end]; + let mut i = 0; + while i < end { + let e = advance_byte_level(text, tags, i, end); emit(Span { - start: start as u32, - end: i as u32, + start: i as u32, + end: e as u32, }); + i = e; } } diff --git a/tokenizers/atomsplit/src/fsm/cl100k.rs b/tokenizers/atomsplit/src/fsm/cl100k.rs index 95f65aca14..5450fd2843 100644 --- a/tokenizers/atomsplit/src/fsm/cl100k.rs +++ b/tokenizers/atomsplit/src/fsm/cl100k.rs @@ -23,17 +23,17 @@ pub fn fsm_cl100k_cap(text: &[u8], tags: &[u8], out: &mut [Span], digit_cap: usi w } -/// The scan under [`fsm_cl100k_cap`]: hands each token to `emit` the moment it is cut, -/// so a caller can consume tokens in place instead of collecting a span buffer first. -pub fn scan_cl100k_cap(text: &[u8], tags: &[u8], digit_cap: usize, mut emit: impl FnMut(Span)) { - debug_assert!(tags.len() >= text.len()); - // Leading-atom values, as `const` so the `match` below is a dense jump table (not an if-cascade): - // the dispatch is O(1) and a token never pays for a rule it can't start (e.g. non-number tokens - // never test the number rule — which is what the POC's const-gating removed by hand; here it's free). - let end = text.len(); - // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior - // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) - let tags = &tags[..end]; +/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of +/// the cl100k-family regex. [`scan_cl100k_cap`] loops it over the whole text; the masked scanner +/// re-derives tokens with it where its batch masks are not trustworthy. +#[inline(always)] +pub(super) fn advance_cl100k_cap( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, + digit_cap: usize, +) -> usize { let letters = |a: usize| run_end(tags, a, end, mask::LETTER); // rule 4 body: `[^\s\p{L}\p{N}]+[\r\n]*` from `sp0` (any leading space already consumed). Returns // the run end, or `sp0` if there is no "other" run there (caller then treats it as whitespace). @@ -49,73 +49,84 @@ pub fn scan_cl100k_cap(text: &[u8], tags: &[u8], digit_cap: usize, mut emit: imp // rules 5-7 (`\s*[\r\n] | \s+(?!\S) | \s+`) → the shared `ws_tail`. let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; - let mut i = 0; - while i < end { - let start = i; - let b = text[i]; - match tags[i] & 0x0F { - // rule 2: `\p{L}+` - LET => i = letters(i), - // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`) - NW | NO => { - let (mut p, mut cnt) = (i, 0); - while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) { - p += char_len(text[p]); - cnt += 1; - } - i = p; - } - // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7 - SPC => { - let a = i + 1; // Space is ASCII (0x20) - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - let p = other(a); - if p > a { p } else { ws(i) } - }; + let b = text[i]; + match tags[i] & 0x0F { + // rule 2: `\p{L}+` + LET => letters(i), + // rule 3: `\p{N}{1,cap}` (cap = 3 cl100k, 1 Qwen2, MAX for `\p{N}+`) + NW | NO => { + let (mut p, mut cnt) = (i, 0); + while p < end && cnt < digit_cap && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; } - // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N) - WSO => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - ws(i) - }; + p + } + // Space: rule 2 (space prefix + `\p{L}+`) | rule 4 (` ` + "other") | rules 5-7 + SPC => { + let a = i + 1; // Space is ASCII (0x20) + if a < end && (tags[a] & 0x0F) == LET { + letters(a) + } else { + let p = other(a); + if p > a { p } else { ws(i) } } - // Newline: whitespace (rule 5 ends at the last newline) - NLN => i = ws(i), - // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4 - APO => { - let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive) - i = if adv > 0 { - i + adv - } else { - let a = i + 1; // Apostrophe is ASCII (0x27) - if a < end && (tags[a] & 0x0F) == LET { - letters(a) - } else { - other(i) - } // c ∈ NOT_WS_L_N ⇒ > i - }; + } + // WsOther: rule 2 (prefix + `\p{L}+`) | whitespace (never rule 4 — not in NOT_WS_L_N) + WSO => { + let a = i + char_len(b); + if a < end && (tags[a] & 0x0F) == LET { + letters(a) + } else { + ws(i) } - // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N): - // rule 2 (prefix + `\p{L}+`) | rule 4 - MRK | CON | PUN | SYM | NMO | CTL => { - let a = i + char_len(b); - i = if a < end && (tags[a] & 0x0F) == LET { + } + // Newline: whitespace (rule 5 ends at the last newline) + NLN => ws(i), + // Apostrophe: rule 1 (contraction) | rule 2 (prefix + `\p{L}+`) | rule 4 + APO => { + let adv = contraction(text, i); // rule 1: `'s 't 're 've 'm 'll 'd` (case-insensitive) + if adv > 0 { + i + adv + } else { + let a = i + 1; // Apostrophe is ASCII (0x27) + if a < end && (tags[a] & 0x0F) == LET { letters(a) } else { other(i) - }; // c ∈ NOT_WS_L_N ⇒ > i + } // c ∈ NOT_WS_L_N ⇒ > i } - // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. - _ => i += char_len(b), } + // Mark | Connector | Punct | SymOther | NumericOther | Control (all in NOT_WS_L_N): + // rule 2 (prefix + `\p{L}+`) | rule 4 + MRK | CON | PUN | SYM | NMO | CTL => { + let a = i + char_len(b); + if a < end && (tags[a] & 0x0F) == LET { + letters(a) + } else { + other(i) + } // c ∈ NOT_WS_L_N ⇒ > i + } + // Sentinel / MultiByte / Cont — never a char-start atom; emit one char defensively. + _ => i + char_len(b), + } +} + +/// The scan under [`fsm_cl100k_cap`]: hands each token to `emit` the moment it is cut, +/// so a caller can consume tokens in place instead of collecting a span buffer first. +pub fn scan_cl100k_cap(text: &[u8], tags: &[u8], digit_cap: usize, mut emit: impl FnMut(Span)) { + debug_assert!(tags.len() >= text.len()); + let end = text.len(); + // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior + // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) + let tags = &tags[..end]; + let mut i = 0; + while i < end { + let e = advance_cl100k_cap(text, tags, i, end, digit_cap); emit(Span { - start: start as u32, - end: i as u32, + start: i as u32, + end: e as u32, }); + i = e; } } diff --git a/tokenizers/atomsplit/src/fsm/deepseek.rs b/tokenizers/atomsplit/src/fsm/deepseek.rs index 75d3131e61..b130ebf0c4 100644 --- a/tokenizers/atomsplit/src/fsm/deepseek.rs +++ b/tokenizers/atomsplit/src/fsm/deepseek.rs @@ -49,6 +49,159 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { w } +/// Maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those), +/// ZWJ/ZWNJ (not `\p{L}∪\p{M}`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`). +/// BYTE-wise (`p += 1`, continuation bytes stay in-run) — the `char_len`-per-char form was ~2× +/// slower (see `run_end`'s note). Hot inner loop of the latin path. +#[inline(always)] +fn letter_run(text: &[u8], tags: &[u8], a: usize, end: usize) -> usize { + let mut p = a; + // ZWJ/ASM are tags (no text peek); only the CJK-range exclusion still peeks text. + while p < end { + let t = tags[p]; + if t == CONT + || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p)) + { + p += 1; + } else { + break; + } + } + p +} + +/// Is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)? +#[inline(always)] +fn is_lm(text: &[u8], tags: &[u8], a: usize, end: usize) -> bool { + a < end + && in_mask(tags[a], mask::LETTER_MARK) + && tags[a] != ASM + && tags[a] != ZWJ + && !ds_is_cjk_at(text, a) +} + +/// Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed); +/// `sp0` if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those, +/// so a CJK punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`). +#[inline(always)] +fn punct(text: &[u8], tags: &[u8], sp0: usize, end: usize) -> usize { + let mut p = sp0; + while p < end && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM) && !ds_is_cjk_at(text, p) + { + p += char_len(text[p]); + } + if p > sp0 { + while p < end && tags[p] == NLN { + p += char_len(text[p]); + } + } + p +} + +/// Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own +/// Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a +/// following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?` +/// prefix. +#[inline(always)] +fn ds_ws(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + let re = run_end(tags, i, end, mask::WS); + let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re)); + if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { + i + r + 1 + } else if re == end || next_isolated { + re // whole ws run is one token + } else { + let mut last = re - 1; + while last > i && text[last] & 0xC0 == 0x80 { + last -= 1; + } + if last > i { last } else { re } + } +} + +/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of +/// the deepseek Sequence. The two multi-emit paths of [`scan_deepseek`] decompose per token: a +/// gap run followed by letters ends at its LAST gap char (the next dispatch, on that char, +/// takes the prefix-plus-letters path), and a CJK run is one same-kind sub-run per call. The +/// masked scanner re-derives tokens with this where its batch masks are not trustworthy. +#[inline(always)] +pub(super) fn advance_deepseek(text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + if ds_is_cjk_at(text, i) { + let is_letter = in_mask(tags[i], mask::LETTER_MARK); + let mut p = i + 3; // CJK-range chars are all 3-byte (leads E3..E9) + while p < end && ds_is_cjk_at(text, p) && in_mask(tags[p], mask::LETTER_MARK) == is_letter { + p += 3; + } + return p; + } + if matches!(tags[i] & 0x0F, NMO | CTL) || tags[i] == ZWJ { + let (mut p, mut last) = (i, i); + while p < end && (matches!(tags[p] & 0x0F, NMO | CTL) || tags[p] == ZWJ) { + last = p; + p += char_len(text[p]); + } + return if is_lm(text, tags, p, end) { + if last > i { + last // gap sans the prefix char + } else { + letter_run(text, tags, p, end) // prefix char + `[\p{L}\p{M}]+` + } + } else { + p + }; + } + let b = text[i]; + match tags[i] & 0x0F { + NW | NO => { + let (mut p, mut cnt) = (i, 0); + while p < end && cnt < 3 && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + p + } + LET | MRK => { + if tags[i] == ASM { + punct(text, tags, i, end) + } else { + letter_run(text, tags, i, end) + } + } + SPC => { + let a = i + 1; // Space is ASCII (0x20) + if is_lm(text, tags, a, end) { + letter_run(text, tags, a, end) + } else if a < end && ds_is_cjk_at(text, a) { + ds_ws(text, tags, i, end) + } else { + let p = punct(text, tags, a, end); + if p > a { p } else { ds_ws(text, tags, i, end) } + } + } + WSO => { + let a = i + char_len(b); + if is_lm(text, tags, a, end) { + letter_run(text, tags, a, end) + } else { + ds_ws(text, tags, i, end) + } + } + NLN => ds_ws(text, tags, i, end), + CON | PUN | APO | SYM => { + if b.is_ascii_punctuation() && i + 1 < end && text[i + 1].is_ascii_alphabetic() { + let mut p = i + 1; + while p < end && text[p].is_ascii_alphabetic() { + p += 1; + } + p + } else { + punct(text, tags, i, end) // c ∈ PUNCT_SYM ⇒ > i + } + } + _ => i + char_len(b), + } +} + /// The scan under [`fsm_deepseek`]: hands each token to `emit` the moment it is cut, /// so a caller can consume tokens in place instead of collecting a span buffer first. pub fn scan_deepseek(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { @@ -61,69 +214,10 @@ pub fn scan_deepseek(text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { // Tie `tags.len() == end` so the optimizer drops the per-byte bounds check on every interior // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) let tags = &tags[..end]; - // maximal `[\p{L}\p{M}]+` run from `a`, stopping at CJK-range chars (Split-2 took those), ZWJ/ZWNJ - // (not `\p{L}∪\p{M}` — see `ds_breaks`), and Other_Alphabetic symbols (`ASM`, categorically `\p{S}`). - // BYTE-wise (`p += 1`, continuation bytes stay in-run, `ds_breaks` only fires at a lead) — the - // `char_len`-per-char form was ~2× slower (see `run_end`'s note). Hot inner loop of the latin path. - let letter_run = |a: usize| -> usize { - let mut p = a; - // ZWJ/ASM are now tags (no text peek); only the CJK-range exclusion still peeks text. - while p < end { - let t = tags[p]; - if t == CONT - || (in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ && !ds_is_cjk_at(text, p)) - { - p += 1; - } else { - break; - } - } - p - }; - // is `text[a]` the start of a deepseek letter/mark char (alt-2 run body / space-prefix target)? - let is_lm = |a: usize| { - a < end - && in_mask(tags[a], mask::LETTER_MARK) - && tags[a] != ASM - && tags[a] != ZWJ - && !ds_is_cjk_at(text, a) - }; - // Split-3 alt-3 tail `[\p{P}\p{S}]+[\r\n]*` from `sp0` (a leading space is already consumed); `sp0` - // if there is no punct/sym run there. STOPS at CJK-range chars — Split-1 isolated those, so a CJK - // punct (・) is never merged into a non-CJK punct run (`!・` → `!`, `・`, not `!・`). - let punct = |sp0: usize| -> usize { - let mut p = sp0; - while p < end - && (in_mask(tags[p], mask::PUNCT_SYM) || tags[p] == ASM) - && !ds_is_cjk_at(text, p) - { - p += char_len(text[p]); - } - if p > sp0 { - while p < end && tags[p] == NLN { - p += char_len(text[p]); - } - } - p - }; - // Split-3 alts d/e/f (whitespace). Unlike cl100k: a ws run FOLLOWED BY a digit/CJK is its own - // Sequence piece (Split-1/2 isolated the next match) → `\s+(?!\S)` takes the WHOLE run; only a - // following letter/punct (same Split-3 piece) leaves the last ws char for its ` ?`/`[^…]?` prefix. - let ws = |i: usize| -> usize { - let re = run_end(tags, i, end, mask::WS); - let next_isolated = re < end && (in_mask(tags[re], mask::NUMBER) || ds_is_cjk_at(text, re)); - if let Some(r) = text[i..re].iter().rposition(|&x| x == 0x0A || x == 0x0D) { - i + r + 1 - } else if re == end || next_isolated { - re // whole ws run is one token - } else { - let mut last = re - 1; - while last > i && text[last] & 0xC0 == 0x80 { - last -= 1; - } - if last > i { last } else { re } - } - }; + let letter_run = |a: usize| letter_run(text, tags, a, end); + let is_lm = |a: usize| is_lm(text, tags, a, end); + let punct = |sp0: usize| punct(text, tags, sp0, end); + let ws = |i: usize| ds_ws(text, tags, i, end); let mut i = 0; while i < end { diff --git a/tokenizers/atomsplit/src/fsm/masked.rs b/tokenizers/atomsplit/src/fsm/masked.rs new file mode 100644 index 0000000000..aa6f1795a9 --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked.rs @@ -0,0 +1,440 @@ +//! Boundary-mask scanners: the SIMD replacement for the per-run scalar walk of the regex-shaped +//! FSMs (one scheme module per regex family, sharing this walker). +//! +//! A scalar FSM advances one token at a time: a rule dispatch, then a per-byte run scan. A +//! masked scanner instead classifies 64 bytes at once into per-class bitmasks (bit k set = +//! "byte k is a letter") and computes every token start in the batch with a few dozen +//! branch-free u64 operations. That works because these regexes are class-run languages with +//! one or two chars of context: a token starts exactly at a class change that is not an +//! absorbed prefix, at a whitespace-run edge, or at a contraction edge, and each of those +//! conditions is a shifted-mask expression. The boundary algebra is transcribed per scheme from +//! gigatoken's `src/pretokenize/fast/` scanners (MIT), with one structural difference: +//! gigatoken classifies raw bytes in-batch and falls back to a per-char loop whenever a batch +//! contains a non-ASCII byte, while these scanners build their masks from the +//! [`crate::classify`] tag stream, which is already SIMD and covers all of Unicode. A +//! continuation byte's tag is [`Atom::Cont`]; the fill step gives it its char's class, so byte +//! adjacency equals char adjacency and the same algebra applies to non-ASCII batches. +//! +//! # Trust boundaries +//! +//! Each scheme's `batch_masks` returns `(boundary, bad)` for one 64-byte batch. A `boundary` +//! bit is a proven token start. A `bad` bit means the algebra cannot decide that byte (batch- +//! edge straddles, char-counted rules over multi-byte chars, run-contextual classes; each +//! scheme documents its own list). `boundary & bad` is always 0, and no span is emitted across +//! a bad zone: the walker re-derives tokens there with the scheme's scalar `advance`, the same +//! rules the plain scan runs, and resumes on masks at the next batch. The scalar scans stay the +//! ground truth; the `masked_*` tests (tests/fsm.rs) pin byte-exactness at every batch-edge +//! offset. +//! +//! # Targets +//! +//! The per-target work is confined to [`block`] (64-byte predicate masks): aarch64 NEON, +//! x86_64 SSE2 and wasm32 simd128 have kernels; every other target delegates the +//! `scan_*_masked` entry points to the scalar scans, so correctness never depends on a SIMD +//! path being present. +//! +//! Inputs must be well-formed UTF-8 (the crate-level contract); the fill step relies on +//! continuation runs of at most 3 bytes. + +use super::*; + +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod block; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod byte_level; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod cl100k; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod deepseek; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod o200k; + +/// [`fsm_byte_level`] over the masked scanner: writes spans into `out` (len >= `text.len()`) +/// and returns the count. +#[must_use] +pub fn fsm_byte_level_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_byte_level_masked(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The masked twin of [`scan_byte_level`]: same tokens, same emit order. Targets without a +/// [`block`] kernel delegate to the scalar scan. +pub fn scan_byte_level_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + walk(&byte_level::ByteLevelMasked, text, tags, emit); + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + scan_byte_level(text, tags, emit); +} + +/// [`fsm_cl100k_cap`] over the masked scanner: writes spans into `out` (len >= `text.len()`) +/// and returns the count. +#[must_use] +pub fn fsm_cl100k_cap_masked( + text: &[u8], + tags: &[u8], + out: &mut [Span], + digit_cap: usize, +) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_cl100k_cap_masked(text, tags, digit_cap, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The masked twin of [`scan_cl100k_cap`]: same tokens, same emit order. Digit caps other than +/// 1, 3 and `usize::MAX` (none ship today) and targets without a [`block`] kernel delegate to +/// the scalar scan. +pub fn scan_cl100k_cap_masked(text: &[u8], tags: &[u8], digit_cap: usize, emit: impl FnMut(Span)) { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + { + if matches!(digit_cap, 1 | 3 | usize::MAX) { + walk(&cl100k::Cl100kMasked { digit_cap }, text, tags, emit); + } else { + scan_cl100k_cap(text, tags, digit_cap, emit); + } + } + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + scan_cl100k_cap(text, tags, digit_cap, emit); +} + +/// [`fsm_o200k`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and +/// returns the count. +#[must_use] +pub fn fsm_o200k_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_o200k_masked(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// [`fsm_tekken`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and +/// returns the count. +#[must_use] +pub fn fsm_tekken_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_tekken_masked(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The masked twin of [`scan_o200k`]: same tokens, same emit order. Targets without a +/// [`block`] kernel delegate to the scalar scan. +pub fn scan_o200k_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + walk(&o200k::O200kMasked::, text, tags, emit); + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + scan_o200k(text, tags, emit); +} + +/// The masked twin of [`scan_tekken`]: same tokens, same emit order. Targets without a +/// [`block`] kernel delegate to the scalar scan. +pub fn scan_tekken_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + walk(&o200k::O200kMasked::, text, tags, emit); + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + scan_tekken(text, tags, emit); +} + +/// [`fsm_deepseek`] over the masked scanner: writes spans into `out` (len >= `text.len()`) and +/// returns the count. +#[must_use] +pub fn fsm_deepseek_masked(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize { + debug_assert!(out.len() >= text.len() && tags.len() >= text.len()); + let mut w = 0usize; + scan_deepseek_masked(text, tags, |span| { + // SAFETY: tokens partition the input, so `w < #tokens <= text.len() <= out.len()`. + unsafe { *out.get_unchecked_mut(w) = span }; + w += 1; + }); + w +} + +/// The masked twin of [`scan_deepseek`]: same tokens, same emit order. Targets without a +/// [`block`] kernel delegate to the scalar scan. +pub fn scan_deepseek_masked(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + walk(&deepseek::DeepSeekMasked, text, tags, emit); + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + scan_deepseek(text, tags, emit); +} + +/// One masked scheme: the batch classifier and the scalar rules the walker falls back on. +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +trait MaskedFsm { + /// `(boundary, bad)` for `text[scan..scan + 64]`: `boundary` bit k = proven token start at + /// `scan + k`, `bad` bit k = the walker must re-derive byte `scan + k` with + /// [`Self::advance`]. `boundary & bad` must be 0. Callers guarantee `scan + 64 < + /// text.len()` (one lookahead tag is readable). + fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64); + + /// Scalar ground truth: end of the token starting at `i` (`i < end`, `i` on a token + /// boundary). + fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize; +} + +/// The batch walker: consume proven token starts batch by batch, re-derive bad zones with the +/// scheme's scalar rules. `pending` is always the start of the open (not yet emitted) token. +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +fn walk(scheme: &impl MaskedFsm, text: &[u8], tags: &[u8], mut emit: impl FnMut(Span)) { + let end = text.len(); + let tags = &tags[..end]; + let mut pending = 0usize; + let mut scan = 0usize; + // Each batch reads the lookahead tag at scan + 64 (the `\s+(?!\S)` bit-63 rules), so the + // last <= 64 bytes always go through the scalar tail below. + while scan + 64 < end { + if scan + 64 <= pending { + // A scalar re-derivation ran past this whole batch. + scan += 64; + continue; + } + let (boundary, mut bad) = scheme.batch_masks(text, tags, scan); + let mut m = boundary; + if pending > scan { + // Bits at or below `pending` are inside or at the start of the open token; they are + // stale leftovers of a scalar re-derivation that entered this batch. + let done = pending - scan + 1; + m = if done >= 64 { + 0 + } else { + m & (u64::MAX << done) + }; + } + // Interleave: consume proven starts below the next bad zone, re-derive the zone with + // the scalar rules, repeat. A span must never be emitted across an unresolved zone, so + // starts above one cannot pair with `pending` from below it. + loop { + let zone = if bad == 0 { + 64 + } else { + bad.trailing_zeros() as usize + }; + while m != 0 { + let j = m.trailing_zeros() as usize; + if j >= zone { + break; + } + let p = scan + j; + if p > pending { + emit(Span { + start: pending as u32, + end: p as u32, + }); + pending = p; + } + m &= m - 1; + } + if bad == 0 { + break; + } + // The zone's contiguous extent; the scalar rules resolve through its end (their + // tokens may overshoot it, or the whole batch — later bits fall to the guards). + let zone_end = zone + ((!(bad >> zone)).trailing_zeros() as usize).min(64 - zone); + while pending < scan + zone_end { + let e = scheme.advance(text, tags, pending, end); + emit(Span { + start: pending as u32, + end: e as u32, + }); + pending = e; + } + bad = if zone_end >= 64 { + 0 + } else { + bad & (u64::MAX << zone_end) + }; + } + scan += 64; + } + while pending < end { + let e = scheme.advance(text, tags, pending, end); + emit(Span { + start: pending as u32, + end: e as u32, + }); + pending = e; + } +} + +// ── shared u64 helpers (platform-independent; the scheme modules compose these) ──────────────── + +/// The two continuation-run masks the fill steps need: `c2` bit k = bytes k and k-1 are both +/// continuations, `c3` = bytes k, k-1, k-2 all are. +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn cont_runs(c: u64) -> (u64, u64) { + let c2 = c & (c << 1); + (c2, c2 & (c << 2)) +} + +/// Fill: every continuation byte of a char whose lead is in `m` joins `m`, so byte adjacency +/// equals char adjacency (UTF-8 chars are at most 4 bytes: 3 hops cover every continuation). +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn fill(m: u64, c: u64, c2: u64, c3: u64) -> u64 { + m | ((m << 1) & c) | ((m << 2) & c2) | ((m << 3) & c3) +} + +/// Smear `seed` upward (toward higher bits) through contiguous set bits of `within`, in log +/// steps (via gigatoken's `cl100k_family.rs`, MIT). +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn smear_up(seed: u64, within: u64) -> u64 { + let mut a = seed; + let mut m = within; + let mut sh = 1u32; + while sh < 64 { + a |= (a << sh) & m; + m &= m << sh; + sh <<= 1; + } + a +} + +/// Token-start bits inside ASCII digit runs for `\p{N}{1,3}`: each run splits into 3-char +/// tokens, so boundaries sit at run start + 3k (via gigatoken's `mask.rs`, MIT). Callers keep +/// multi-byte digit chars out of `d`: their grouping is char-counted, and byte hops would +/// misphase it. +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn digit_run_splits3(d: u64) -> u64 { + let mut b = d & !(d << 1); // run starts + // A start at p re-arms at p+3 while the run continues: hop condition c = "p..p+3 all + // digits". Log-doubling covers 64-bit runs in 5 steps. + let mut c = d & (d >> 1) & (d >> 2) & (d >> 3); + let mut sh = 3u32; + while sh < 64 { + b |= (b & c) << sh; + c &= c >> sh; + sh <<= 1; + } + b +} + +/// `x << n`, saturating to 0 at `n >= 64` (`trailing_zeros` on an empty mask yields 64). +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn shl_sat(x: u64, n: u32) -> u64 { + if n >= 64 { 0 } else { x << n } +} + +/// The lead of the char containing `tags[p]`: walk back over continuation tags (at most 3 on +/// well-formed UTF-8; `p` itself may already be the lead). +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +#[inline(always)] +fn char_lead(tags: &[u8], mut p: usize) -> usize { + while p > 0 && tags[p] & 0x0F == CONT { + p -= 1; + } + p +} diff --git a/tokenizers/atomsplit/src/fsm/masked/block.rs b/tokenizers/atomsplit/src/fsm/masked/block.rs new file mode 100644 index 0000000000..3e06e7c898 --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked/block.rs @@ -0,0 +1,427 @@ +//! Per-target 64-byte block classifiers for the masked scanners. A block loads 64 consecutive +//! bytes into four SIMD registers once; each method then answers one per-byte predicate for the +//! whole block as a u64 bitmask (bit k = byte k passes). Everything downstream of these masks is +//! platform-independent u64 arithmetic in the scheme modules, so this file is the entire +//! per-target surface. +//! +//! Three targets have a kernel: aarch64 (NEON, baseline), x86_64 (SSE2, baseline) and wasm32 +//! with `simd128`. Any other target never reaches this module — the `scan_*_masked` entry +//! points delegate to the scalar scans there. + +#![allow(dead_code)] // arch-gated: each build compiles one target's kernel, and schemes land +// one by one, so some predicates are unused until their scheme arrives. + +// ── aarch64 / NEON ────────────────────────────────────────────────────────────────────────────── +#[cfg(target_arch = "aarch64")] +pub(crate) use neon::Block; + +#[cfg(target_arch = "aarch64")] +mod neon { + use core::arch::aarch64::*; + + /// 64 bytes in four NEON registers. `tag` methods fold the refinement nibble away first + /// (`& 0x0F`), `full`/byte methods compare the raw byte. + pub(crate) struct Block { + v: [uint8x16_t; 4], + } + + /// simdjson's arm64 movemask: 4 mask vectors (64 lanes of 0x00/0xFF) to one u64, bit i = + /// lane i. The 4-`addp` reduction is pinned as asm (via gigatoken's `mask.rs`, MIT): written + /// with `vpaddq_u8`, LLVM rewrites the pairwise adds into uzp1/uzp2/orr triples and the call + /// grows from 9 to 17 vector ops. + #[inline(always)] + unsafe fn movemask64(v0: uint8x16_t, v1: uint8x16_t, v2: uint8x16_t, v3: uint8x16_t) -> u64 { + // SAFETY: pure NEON register arithmetic, no memory access beyond the 16-byte constant. + unsafe { + const W: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; + let w = vld1q_u8(W.as_ptr()); + let mut a0 = vandq_u8(v0, w); + let a1 = vandq_u8(v1, w); + let a2 = vandq_u8(v2, w); + let a3 = vandq_u8(v3, w); + core::arch::asm!( + "addp {a0:v}.16b, {a0:v}.16b, {a1:v}.16b", + "addp {a2:v}.16b, {a2:v}.16b, {a3:v}.16b", + "addp {a0:v}.16b, {a0:v}.16b, {a2:v}.16b", + "addp {a0:v}.16b, {a0:v}.16b, {a0:v}.16b", + a0 = inout(vreg) a0, + a1 = in(vreg) a1, + a2 = inout(vreg) a2 => _, + a3 = in(vreg) a3, + options(pure, nomem, nostack, preserves_flags), + ); + vgetq_lane_u64::<0>(vreinterpretq_u64_u8(a0)) + } + } + + impl Block { + /// Load `bytes[at..at + 64]`. + /// + /// # Safety + /// + /// `at + 64 <= bytes.len()`. NEON loads are alignment-free. + #[inline(always)] + pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self { + debug_assert!(at + 64 <= bytes.len()); + // SAFETY: the fn contract puts all four 16-byte loads in bounds. + unsafe { + let p = bytes.as_ptr().add(at); + Self { + v: [ + vld1q_u8(p), + vld1q_u8(p.add(16)), + vld1q_u8(p.add(32)), + vld1q_u8(p.add(48)), + ], + } + } + } + + #[inline(always)] + fn mask(&self, f: impl Fn(uint8x16_t) -> uint8x16_t) -> u64 { + // SAFETY: register arithmetic only. + unsafe { movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3])) } + } + + #[inline(always)] + fn any(&self, f: impl Fn(uint8x16_t) -> uint8x16_t) -> bool { + // SAFETY: register arithmetic only. + unsafe { + let o = vorrq_u8( + vorrq_u8(f(self.v[0]), f(self.v[1])), + vorrq_u8(f(self.v[2]), f(self.v[3])), + ); + vmaxvq_u8(o) != 0 + } + } + + /// Bytes whose low nibble equals `k`. + #[inline(always)] + pub(crate) fn eq_tag(&self, k: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { vceqq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(k)) }) + } + + /// Bytes in `lo..=lo + span` (the raw byte — text blocks). + #[inline(always)] + pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { vcleq_u8(vsubq_u8(v, vdupq_n_u8(lo)), vdupq_n_u8(span)) }) + } + + /// Bytes whose low nibble is in `lo..=lo + span`. + #[inline(always)] + pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + vcleq_u8( + vsubq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(lo)), + vdupq_n_u8(span), + ) + }) + } + + /// Bytes equal to `k` (the raw byte — refined tags, or text bytes). + #[inline(always)] + pub(crate) fn eq_full(&self, k: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { vceqq_u8(v, vdupq_n_u8(k)) }) + } + + /// Any byte with low nibble in `lo..=lo + span`? (Cheaper than `range_tag` when only + /// presence matters — no movemask.) + #[inline(always)] + pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool { + // SAFETY: register arithmetic only. + self.any(|v| unsafe { + vcleq_u8( + vsubq_u8(vandq_u8(v, vdupq_n_u8(0x0F)), vdupq_n_u8(lo)), + vdupq_n_u8(span), + ) + }) + } + + /// Any byte equal to `k`? + #[inline(always)] + pub(crate) fn any_eq_full(&self, k: u8) -> bool { + // SAFETY: register arithmetic only. + self.any(|v| unsafe { vceqq_u8(v, vdupq_n_u8(k)) }) + } + + /// ASCII letters `[A-Za-z]` (text blocks). + #[inline(always)] + pub(crate) fn ascii_alpha(&self) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + vcleq_u8( + vsubq_u8(vorrq_u8(v, vdupq_n_u8(0x20)), vdupq_n_u8(b'a')), + vdupq_n_u8(25), + ) + }) + } + + /// ASCII punctuation (the four `is_ascii_punctuation` ranges, text blocks). + #[inline(always)] + pub(crate) fn ascii_punct(&self) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + let r = |v: uint8x16_t, lo: u8, span: u8| { + vcleq_u8(vsubq_u8(v, vdupq_n_u8(lo)), vdupq_n_u8(span)) + }; + vorrq_u8( + vorrq_u8(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)), + vorrq_u8(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)), + ) + }) + } + } +} + +// ── x86_64 / SSE2 ────────────────────────────────────────────────────────────────────────────── +#[cfg(target_arch = "x86_64")] +pub(crate) use sse2::Block; + +#[cfg(target_arch = "x86_64")] +mod sse2 { + use core::arch::x86_64::*; + + /// 64 bytes in four SSE2 registers; `pmovmskb` is the native movemask, 16 bits per register. + /// SSE2 is baseline on x86_64, so no runtime detection is needed. (SSE2 has no unsigned + /// byte compare: `x <= span` is done as `max(x, span) == span`.) + pub(crate) struct Block { + v: [__m128i; 4], + } + + #[inline(always)] + fn movemask64(v0: __m128i, v1: __m128i, v2: __m128i, v3: __m128i) -> u64 { + // SAFETY: register arithmetic only; SSE2 is baseline on x86_64. + unsafe { + (_mm_movemask_epi8(v0) as u16 as u64) + | ((_mm_movemask_epi8(v1) as u16 as u64) << 16) + | ((_mm_movemask_epi8(v2) as u16 as u64) << 32) + | ((_mm_movemask_epi8(v3) as u16 as u64) << 48) + } + } + + #[inline(always)] + fn le(x: __m128i, span: u8) -> __m128i { + // SAFETY: register arithmetic only. + unsafe { + let s = _mm_set1_epi8(span as i8); + _mm_cmpeq_epi8(_mm_max_epu8(x, s), s) + } + } + + impl Block { + /// Load `bytes[at..at + 64]`. + /// + /// # Safety + /// + /// `at + 64 <= bytes.len()`. Unaligned loads (`loadu`). + #[inline(always)] + pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self { + debug_assert!(at + 64 <= bytes.len()); + // SAFETY: the fn contract puts all four 16-byte loads in bounds. + unsafe { + let p = bytes.as_ptr().add(at); + Self { + v: [ + _mm_loadu_si128(p.cast()), + _mm_loadu_si128(p.add(16).cast()), + _mm_loadu_si128(p.add(32).cast()), + _mm_loadu_si128(p.add(48).cast()), + ], + } + } + } + + #[inline(always)] + fn mask(&self, f: impl Fn(__m128i) -> __m128i) -> u64 { + movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3])) + } + + #[inline(always)] + pub(crate) fn eq_tag(&self, k: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + _mm_cmpeq_epi8( + _mm_and_si128(v, _mm_set1_epi8(0x0F)), + _mm_set1_epi8(k as i8), + ) + }) + } + + #[inline(always)] + pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + le( + _mm_sub_epi8( + _mm_and_si128(v, _mm_set1_epi8(0x0F)), + _mm_set1_epi8(lo as i8), + ), + span, + ) + }) + } + + #[inline(always)] + pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { le(_mm_sub_epi8(v, _mm_set1_epi8(lo as i8)), span) }) + } + + #[inline(always)] + pub(crate) fn eq_full(&self, k: u8) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { _mm_cmpeq_epi8(v, _mm_set1_epi8(k as i8)) }) + } + + #[inline(always)] + pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool { + self.range_tag(lo, span) != 0 + } + + #[inline(always)] + pub(crate) fn any_eq_full(&self, k: u8) -> bool { + self.eq_full(k) != 0 + } + + #[inline(always)] + pub(crate) fn ascii_alpha(&self) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + le( + _mm_sub_epi8( + _mm_or_si128(v, _mm_set1_epi8(0x20)), + _mm_set1_epi8(b'a' as i8), + ), + 25, + ) + }) + } + + #[inline(always)] + pub(crate) fn ascii_punct(&self) -> u64 { + // SAFETY: register arithmetic only. + self.mask(|v| unsafe { + let r = |v: __m128i, lo: u8, span: u8| { + le(_mm_sub_epi8(v, _mm_set1_epi8(lo as i8)), span) + }; + _mm_or_si128( + _mm_or_si128(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)), + _mm_or_si128(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)), + ) + }) + } + } +} + +// ── wasm32 / SIMD128 ──────────────────────────────────────────────────────────────────────────── +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +pub(crate) use wasm::Block; + +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +mod wasm { + use core::arch::wasm32::*; + + /// 64 bytes in four v128 registers; `u8x16_bitmask` is the native movemask. + pub(crate) struct Block { + v: [v128; 4], + } + + #[inline(always)] + fn movemask64(v0: v128, v1: v128, v2: v128, v3: v128) -> u64 { + (u8x16_bitmask(v0) as u64) + | ((u8x16_bitmask(v1) as u64) << 16) + | ((u8x16_bitmask(v2) as u64) << 32) + | ((u8x16_bitmask(v3) as u64) << 48) + } + + impl Block { + /// Load `bytes[at..at + 64]`. + /// + /// # Safety + /// + /// `at + 64 <= bytes.len()`. `v128_load` is alignment-free. + #[inline(always)] + pub(crate) unsafe fn load(bytes: &[u8], at: usize) -> Self { + debug_assert!(at + 64 <= bytes.len()); + // SAFETY: the fn contract puts all four 16-byte loads in bounds. + unsafe { + let p = bytes.as_ptr().add(at); + Self { + v: [ + v128_load(p.cast()), + v128_load(p.add(16).cast()), + v128_load(p.add(32).cast()), + v128_load(p.add(48).cast()), + ], + } + } + } + + #[inline(always)] + fn mask(&self, f: impl Fn(v128) -> v128) -> u64 { + movemask64(f(self.v[0]), f(self.v[1]), f(self.v[2]), f(self.v[3])) + } + + #[inline(always)] + pub(crate) fn eq_tag(&self, k: u8) -> u64 { + self.mask(|v| u8x16_eq(v128_and(v, u8x16_splat(0x0F)), u8x16_splat(k))) + } + + #[inline(always)] + pub(crate) fn range_tag(&self, lo: u8, span: u8) -> u64 { + self.mask(|v| { + u8x16_le( + u8x16_sub(v128_and(v, u8x16_splat(0x0F)), u8x16_splat(lo)), + u8x16_splat(span), + ) + }) + } + + #[inline(always)] + pub(crate) fn range_full(&self, lo: u8, span: u8) -> u64 { + self.mask(|v| u8x16_le(u8x16_sub(v, u8x16_splat(lo)), u8x16_splat(span))) + } + + #[inline(always)] + pub(crate) fn eq_full(&self, k: u8) -> u64 { + self.mask(|v| u8x16_eq(v, u8x16_splat(k))) + } + + #[inline(always)] + pub(crate) fn any_range_tag(&self, lo: u8, span: u8) -> bool { + self.range_tag(lo, span) != 0 + } + + #[inline(always)] + pub(crate) fn any_eq_full(&self, k: u8) -> bool { + self.eq_full(k) != 0 + } + + #[inline(always)] + pub(crate) fn ascii_alpha(&self) -> u64 { + self.mask(|v| { + u8x16_le( + u8x16_sub(v128_or(v, u8x16_splat(0x20)), u8x16_splat(b'a')), + u8x16_splat(25), + ) + }) + } + + #[inline(always)] + pub(crate) fn ascii_punct(&self) -> u64 { + self.mask(|v| { + let r = |v: v128, lo: u8, span: u8| { + u8x16_le(u8x16_sub(v, u8x16_splat(lo)), u8x16_splat(span)) + }; + v128_or( + v128_or(r(v, 0x21, 0x0E), r(v, 0x3A, 0x06)), + v128_or(r(v, 0x5B, 0x05), r(v, 0x7B, 0x03)), + ) + }) + } + } +} diff --git a/tokenizers/atomsplit/src/fsm/masked/byte_level.rs b/tokenizers/atomsplit/src/fsm/masked/byte_level.rs new file mode 100644 index 0000000000..9008854341 --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked/byte_level.rs @@ -0,0 +1,146 @@ +//! Masked scheme for the byte-level (GPT-2) regex — the r50k boundary algebra (via gigatoken's +//! `r50k.rs`, MIT) over tag-fed class masks. +//! +//! Every byte-level rule is local: a token starts exactly at a class change that is not a space +//! prefix, at the first byte of a whitespace run, at the last whitespace byte before a +//! non-whitespace (the `\s+(?!\S)` give-back), or at a contraction edge. Bad zones: a +//! whitespace char straddling the batch edge (its give-back is per char, not per byte), an +//! apostrophe too close to the edge for the contraction peek, a multi-byte whitespace char, or +//! a `Sentinel`/`MultiByte` tag. + +use super::super::byte_level::advance_byte_level; +use super::block::Block; +use super::{MaskedFsm, char_lead, cont_runs, fill}; +use crate::fsm::{APO, CONT, LET, NLN, NO, NW, SPC, WSO, in_mask, mask}; + +pub(super) struct ByteLevelMasked; + +impl MaskedFsm for ByteLevelMasked { + #[inline(always)] + fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + batch_masks(text, tags, scan) + } + + #[inline(always)] + fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + advance_byte_level(text, tags, i, end) + } +} + +#[inline(always)] +fn batch_masks(text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + debug_assert!(scan + 64 < tags.len() && tags.len() == text.len()); + // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract. + let b = unsafe { Block::load(tags, scan) }; + if b.any_range_tag(13, 1) { + // Sentinel / MultiByte: the scalar dispatch has a defensive arm for these; keep its + // behavior by refusing the whole batch. + return (0, u64::MAX); + } + let l0 = b.eq_tag(LET); + let d0 = b.range_tag(NW, 1); + let ws0 = b.range_tag(NLN, 2); + let s = b.eq_tag(SPC); + // Apostrophes and continuations only matter to the fixups below; skip their movemask when + // an any-test says the batch has none. + let ap = if b.any_eq_full(APO) { + b.eq_full(APO) + } else { + 0 + }; + let c = if b.any_eq_full(CONT) { + b.eq_full(CONT) + } else { + 0 + }; + let (mut l, mut d, mut ws) = (l0, d0, ws0); + + // Carries: the classes of the byte just before the batch, which is the class of the char + // containing it. + let (pl, pd, pws, ps, po) = if scan == 0 { + (0u64, 0u64, 0u64, 0u64, 0u64) + } else { + match tags[char_lead(tags, scan - 1)] & 0x0F { + LET => (1, 0, 0, 0, 0), + NW | NO => (0, 1, 0, 0, 0), + SPC => (0, 0, 1, 1, 0), + NLN | WSO => (0, 0, 1, 0, 0), + _ => (0, 0, 0, 0, 1), + } + }; + + if c != 0 { + // A char straddling into the batch has its leading continuation bytes take the carry + // class ("other" needs no action: it is derived as the complement). + if c & 1 != 0 { + let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1); + l |= lead_in * pl; + d |= lead_in * pd; + ws |= lead_in * pws; + } + let (c2, c3) = cont_runs(c); + l = fill(l, c, c2, c3); + d = fill(d, c, c2, c3); + ws = fill(ws, c, c2, c3); + if ws & c != 0 { + // Multi-byte whitespace: the `\s+(?!\S)` give-back is one char, not one byte, and + // the algebra below works in bytes. Rare (NBSP and friends); scalar batch. + return (0, u64::MAX); + } + } + let o = !(l | d | ws); + + // The r50k boundary algebra (gigatoken, MIT). A byte starts a token when it is not + // whitespace, does not continue a same-class run, and does not follow a space (the ` ?` + // prefix glues it to the space instead). + let cont_same = (l & ((l << 1) | pl)) | (d & ((d << 1) | pd)) | (o & ((o << 1) | po)); + let after_sp = (s << 1) | ps; + let nb = !ws & !cont_same & !after_sp; + + let mut bad = 0u64; + + // Whitespace-run splits. `split_ok` = the last whitespace byte before a non-whitespace (the + // `\s+(?!\S)` give-back starts a token there); bit 63 needs the lookahead tag. + let mut split_ok = ws & (!ws >> 1); + let la = tags[scan + 64] & 0x0F; + if la == CONT { + // The char at the batch edge straddles out. Whitespace is the only class whose rules + // look at char ends, so only a whitespace lead poisons the tail. + let p = char_lead(tags, scan + 63); + if in_mask(tags[p], mask::WS) { + bad |= u64::MAX << (p - scan); + } + } else if !in_mask(la, mask::WS) { + split_ok |= ws & (1u64 << 63); + } + let pwsb = (ws << 1) | pws; + let wsboundary = ws & (!pwsb | split_ok); + let mut boundary = nb | wsboundary; + + // Contraction fixup, only when the batch has an apostrophe that starts a token: a match + // (case-sensitive, as in the scalar dispatch) absorbs the next 1-2 letters and re-opens a + // token after them. Too close to the edge (i >= 61) the peek and the re-opened bit can + // cross the batch; refuse the tail instead. + if ap != 0 { + let mut cand = ap & boundary; + while cand != 0 { + let i = cand.trailing_zeros() as usize; + cand &= cand - 1; + if i >= 61 { + bad |= u64::MAX << i; + break; + } + let k = match text[scan + i + 1] { + b's' | b't' | b'm' | b'd' => 2, + b'r' | b'v' if text[scan + i + 2] == b'e' => 3, + b'l' if text[scan + i + 2] == b'l' => 3, + _ => 0, + }; + if k != 0 { + boundary &= !(1u64 << (i + 1)); + boundary |= 1u64 << (i + k); + } + } + } + (boundary & !bad, bad) +} diff --git a/tokenizers/atomsplit/src/fsm/masked/cl100k.rs b/tokenizers/atomsplit/src/fsm/masked/cl100k.rs new file mode 100644 index 0000000000..438835073e --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked/cl100k.rs @@ -0,0 +1,313 @@ +//! Masked scheme for the cl100k regex family (cl100k / Llama-3 / GLM at digit cap 3, Qwen2 at +//! cap 1, unbounded `\p{N}+` at `usize::MAX`) — the boundary algebra via gigatoken's +//! `cl100k_family.rs` (MIT) over tag-fed class masks. +//! +//! Boundary rules (the regex is `'(?i:contractions)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,cap}| +//! ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+`): +//! - A letter starts a token unless it continues a letter run, follows a space or a non-newline +//! whitespace char (those always sit at a boundary before a non-ws char and absorb one letter +//! run via the `[^\r\n\p{L}\p{N}]?` prefix), or follows a punct char that is itself at a +//! boundary, i.e. whose own predecessor is neither punct nor space (a two-chars-back test, +//! made char-aware by shifting per the previous char's byte length). +//! - Digits split every `cap` chars from each run start and never absorb a preceding space. +//! - A punct char starts a token unless it continues a punct run or follows a space. +//! - Newlines directly after a punct run are absorbed (`[\r\n]*`). +//! - A whitespace run containing newlines emits one token through its LAST newline +//! (`\s*[\r\n]`), then the give-back rules; NL-free runs split before their last char when +//! followed by non-ws (`\s+(?!\S)`). +//! +//! Bad zones: multi-byte `\p{N}` chars under cap 3 (the grouping is char-counted, byte hops +//! would misphase it) and any digit run whose phase starts before the batch; whitespace runs +//! touching the batch end while the next char is whitespace (their last newline may lie +//! beyond); apostrophes near the batch edge or before a non-ASCII char (`(?i:'s)` also matches +//! `'ſ`); `Sentinel`/`MultiByte` tags. + +use super::super::cl100k::advance_cl100k_cap; +use super::block::Block; +use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up}; +use crate::fsm::{APO, CONT, LET, NLN, NO, NW, SPC, WSO, in_mask, mask}; + +pub(super) struct Cl100kMasked { + /// 1, 3 or `usize::MAX`; the entry point routes any other cap to the scalar scan. + pub(super) digit_cap: usize, +} + +impl MaskedFsm for Cl100kMasked { + #[inline(always)] + fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + batch_masks(text, tags, scan, self.digit_cap) + } + + #[inline(always)] + fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + advance_cl100k_cap(text, tags, i, end, self.digit_cap) + } +} + +/// Boundary carries from the two chars before the batch: P1 is the char containing byte +/// `scan - 1`, P2 the one before it (the two-chars-back absorb test). +#[derive(Default)] +struct Carries { + pl: u64, + ps: u64, + pwt: u64, + po: u64, + pws: u64, + pd: u64, + /// P2 is punct-or-space, for a char lead at bit 0 (P1 entirely before the batch). + c2_os: u64, + /// The same test positioned at the first lead after a P1 that straddles into the batch + /// (P1's own predecessor is then P2). + b2b_in: u64, +} + +fn carries(tags: &[u8], scan: usize, c: u64) -> Carries { + let mut cr = Carries::default(); + if scan == 0 { + return cr; + } + let p1 = char_lead(tags, scan - 1); + let c2v = if p1 == 0 { + 0 + } else { + let t2 = tags[char_lead(tags, p1 - 1)] & 0x0F; + u64::from(t2 == SPC || in_mask(t2, mask::NOT_WS_L_N)) + }; + if c & 1 != 0 { + cr.b2b_in = c2v << (!c).trailing_zeros(); + } else { + cr.c2_os = c2v; + } + match tags[p1] & 0x0F { + LET => cr.pl = 1, + NW | NO => cr.pd = 1, + SPC => { + cr.pws = 1; + cr.ps = 1; + } + WSO => { + cr.pws = 1; + cr.pwt = 1; + } + // A newline is whitespace but never a `[^\r\n\p{L}\p{N}]?` prefix, so pwt stays 0. + NLN => cr.pws = 1, + _ => cr.po = 1, + } + cr +} + +fn batch_masks(text: &[u8], tags: &[u8], scan: usize, digit_cap: usize) -> (u64, u64) { + debug_assert!(scan + 64 < tags.len() && tags.len() == text.len()); + // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract. + let blk = unsafe { Block::load(tags, scan) }; + if blk.any_range_tag(13, 1) { + // Sentinel / MultiByte: the scalar dispatch has a defensive arm for these. + return (0, u64::MAX); + } + let l0 = blk.eq_tag(LET); + let d0 = blk.range_tag(NW, 1); + let nl = blk.eq_tag(NLN); + let s = blk.eq_tag(SPC); + let wt0 = blk.eq_tag(WSO); + let ap = if blk.any_eq_full(APO) { + blk.eq_full(APO) + } else { + 0 + }; + let c = if blk.any_eq_full(CONT) { + blk.eq_full(CONT) + } else { + 0 + }; + + let cr = carries(tags, scan, c); + + // Fill: continuation bytes join their char's class; a char straddling into the batch has + // its leading continuation bytes take P1's class (punct needs no action: it is derived as + // the complement). + let (mut l, mut d, mut wt) = (l0, d0, wt0); + let (c2, c3) = cont_runs(c); + if c != 0 { + if c & 1 != 0 { + let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1); + l |= lead_in * cr.pl; + d |= lead_in * cr.pd; + wt |= lead_in * cr.pwt; + } + l = fill(l, c, c2, c3); + d = fill(d, c, c2, c3); + wt = fill(wt, c, c2, c3); + } + // Per-length char leads (`\s` chars are at most 3 bytes; letters up to 4). + let lead = !c; + let len1 = lead & !(c >> 1); + let len2 = lead & (c >> 1) & !(c >> 2); + let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3); + let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3); + let w2 = wt0 & len2; + let w3 = wt0 & len3; + + let ws_f = s | nl | wt; + let o = !(l | d | ws_f); + + // --- Letters: `[^\r\n\p{L}\p{N}]?\p{L}+` -------------------------------------------------- + // b2back: "the char two back is punct or space", evaluated at each char's lead by shifting + // the prev-byte test by the PREVIOUS char's byte length. + let c_test = ((o | s) << 1) | cr.po | cr.ps; + let b2back = ((c_test & len1) << 1) + | ((c_test & len2) << 2) + | ((c_test & len3) << 3) + | ((c_test & len4) << 4) + | cr.c2_os + | cr.b2b_in; + let p_l = (l << 1) | cr.pl; + let p_s = (s << 1) | cr.ps; + let p_wt = (wt << 1) | cr.pwt; + let p_o = (o << 1) | cr.po; + let absorb = p_o & !b2back; + let b_letters = l0 & !p_l & !p_s & !p_wt & !absorb; + + // --- Digits: `\p{N}{1,cap}` ---------------------------------------------------------------- + // Only 1-byte digit chars can be split by byte hops; multi-byte `\p{N}` runs go to the + // scalar path under cap 3 (bad below). Cap 1 tokens are single chars (any width), and the + // unbounded cap only needs run starts. + let d_ascii = d0 & len1; + let dmb = (d & c) | (d0 & !len1); + let b_digits = match digit_cap { + 3 => { + if d_ascii & (d_ascii >> 1) != 0 { + digit_run_splits3(d_ascii) + } else { + d_ascii + } + } + 1 => d0, + _ => d0 & !((d << 1) | cr.pd), + }; + + // --- Punct: ` ?[^\s\p{L}\p{N}]+[\r\n]*` ---------------------------------------------------- + let b_punct = (o & lead) & !p_o & !p_s; + + // Newlines directly after a punct run are absorbed (`[\r\n]*`). + let abs_seed = nl & ((o << 1) | cr.po); + let abs_n = if abs_seed == 0 { + 0 + } else { + smear_up(abs_seed, nl) + }; + let ws_eff = ws_f & !abs_n; + + let mut bad = if digit_cap == 3 { + dmb | dmb << 1 | dmb >> 1 + } else { + 0 + }; + + // Byte-64 lookahead: is the char at the next batch's first byte non-ws? Decides whether + // ws-like runs touching bit 63 resolve in-batch. A CONT lookahead means a char straddles + // out; treating that as "ws" is safe: a live give-back at bit 63 would need the next char's + // lead at byte 64, which contradicts the straddle. + let la = tags[scan + 64] & 0x0F; + let nn64 = la != CONT && !in_mask(la, mask::WS); + let nn64m = u64::from(nn64).wrapping_neg(); + + // An absorbed newline touching the batch end: if byte 64 is ws the token may continue with + // another newline, and the next batch cannot tell an absorbed `\n` before its bit 0 from a + // ws-run `\n` — defer. + if abs_n >> 63 != 0 && !nn64 { + bad |= 1u64 << 63; + } + + // A ws run touching the batch end resolves in-batch only when byte 64's char is non-ws + // (its last newline and `(?!\S)` split are then all visible); otherwise defer it. + let nonws = !ws_eff; + if ws_eff >> 63 != 0 && !nn64 { + if nonws == 0 { + return (0, u64::MAX); // whole batch one ws run + } + let h = 63 - nonws.leading_zeros(); + bad |= u64::MAX << (h + 1); + } + + // A digit run whose grouping phase did not start inside this batch (a continuation from + // before it, or following a bad zone that may hold digit chars) defers too. + if digit_cap == 3 { + let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd); + if seed != 0 { + bad |= smear_up(seed, d_ascii); + } + } + + // --- Whitespace --------------------------------------------------------------------------- + // Base rule (correct for NL-free runs; NL runs are overridden below): run start, or split + // before the last char when followed by non-ws. + let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff; + let ws_leads = (ws_leads1 | w2 | w3) & !abs_n; + let p_ws = (ws_eff << 1) | cr.pws; + let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61)); + let split_ok = (ws_leads1 & (nonws >> 1)) + | (w2 & (nonws >> 2)) + | (w3 & (nonws >> 3)) + | (edge_last & nn64m); + let mut b_ws = ws_leads & (!p_ws | split_ok); + + // Override every run containing a (non-absorbed) newline: one token through the run's last + // newline, then the give-back rules on the remainder. + let mut runs_n = nl & ws_eff & !bad; + while runs_n != 0 { + let f = runs_n.trailing_zeros(); + let below_gap = nonws & ((1u64 << f) - 1); + let a = if below_gap == 0 { + 0 + } else { + 64 - below_gap.leading_zeros() + }; + let e = (nonws & (u64::MAX << f)).trailing_zeros(); + let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e); + b_ws &= !run_mask; + b_ws |= 1u64 << a; + let q = 63 - (nl & run_mask).leading_zeros(); // last newline in the run + if q + 1 < e { + // Tail after the last newline: starts a token, and its last char splits off before + // the following non-ws char. + b_ws |= 1u64 << (q + 1); + let tail = run_mask & (u64::MAX << (q + 1)); + let tail_leads = ws_leads & tail; + b_ws |= 1u64 << (63 - tail_leads.leading_zeros()); + } + runs_n &= !run_mask; + } + + let mut boundary = b_letters | b_digits | b_punct | b_ws; + + // --- Contractions: `'(?i:[sdmt]|ll|ve|re)` ------------------------------------------------- + let mut cand = ap & boundary & !bad; + while cand != 0 { + let i = cand.trailing_zeros() as usize; + cand &= cand - 1; + if i >= 61 { + bad |= u64::MAX << i; + break; + } + let b1 = text[scan + i + 1]; + if b1 >= 0x80 { + // `(?i:'s)` also matches 'ſ (U+017F): an apostrophe before any non-ASCII char + // defers to the scalar path. + bad |= 0b111u64 << i; + continue; + } + let k = match b1 | 0x20 { + b's' | b'd' | b'm' | b't' => 2, + b'l' if text[scan + i + 2] | 0x20 == b'l' => 3, + b'v' if text[scan + i + 2] | 0x20 == b'e' => 3, + b'r' if text[scan + i + 2] | 0x20 == b'e' => 3, + _ => 0, + }; + if k != 0 { + boundary &= !(1u64 << (i + 1)); + boundary |= 1u64 << (i + k); + } + } + + (boundary & !bad, bad) +} diff --git a/tokenizers/atomsplit/src/fsm/masked/deepseek.rs b/tokenizers/atomsplit/src/fsm/masked/deepseek.rs new file mode 100644 index 0000000000..4598d4ae71 --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked/deepseek.rs @@ -0,0 +1,306 @@ +//! Masked scheme for the deepseek-v3 Sequence (digits `\p{N}{1,3}` → CJK-range runs → the big +//! regex) — boundary rules derived from [`scan_deepseek`]'s dispatch, following gigatoken's +//! `deepseek_v3.rs` scoping (MIT). +//! +//! Deepseek has no case rules and no contractions, but three shapes of its own: +//! - The alt-2 prefix class is `[^\r\n\p{L}\p{P}\p{S}]?`: letters absorb a preceding space, +//! non-newline whitespace char, or the LAST char of a gap run (Control / NumericOther / +//! ZWJ — chars matching no alternative), and never a punct char or a digit. +//! - alt-1 `[ascii-punct][A-Za-z]+`: an ASCII punct char at a token start absorbs a following +//! ASCII-letter run. Where that run collides with a non-ASCII letter or mark, the two rules +//! diverge (`[A-Za-z]+` stops, `[\p{L}\p{M}]+` would continue); those collision bits defer. +//! - A whitespace run followed by a digit or CJK char is one whole token (Split-1/2 isolated +//! the follower), so the `\s+(?!\S)` give-back is gated on the follower's class. +//! +//! CJK-range chars are closed units re-split into same-kind sub-runs, and every other rule +//! stops at them; a batch containing any byte with a lead in `0xE3..=0xE9` (a superset of the +//! CJK ranges) defers whole, before the tag masks are built. On CJK-dominated text the scanner +//! is then the scalar scan plus one 64-byte test per batch, which is the accepted trade: the +//! win is on latin/code text, and CJK batches resolve through the same scalar rules as before. + +use super::super::deepseek::advance_deepseek; +use super::block::Block; +use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up}; +use crate::fsm::{ASM, CON, CONT, CTL, LET, MRK, NLN, NMO, NO, NW, SPC, WSO, ZWJ, in_mask, mask}; + +pub(super) struct DeepSeekMasked; + +impl MaskedFsm for DeepSeekMasked { + #[inline(always)] + fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + batch_masks(text, tags, scan) + } + + #[inline(always)] + fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + advance_deepseek(text, tags, i, end) + } +} + +/// A deepseek letter-run member tag: `[\p{L}\p{M}]` minus the refined Marks (ASM/ZWJ). The +/// CJK-range exclusion is handled by the bad cover, not here. +#[inline(always)] +fn is_member(t: u8) -> bool { + in_mask(t, mask::LETTER_MARK) && t != ASM && t != ZWJ +} + +#[derive(Default)] +struct Carries { + pl: u64, + ps: u64, + pwt: u64, + po: u64, + pws: u64, + pd: u64, + pgap: u64, + /// Byte `scan - 1` is an ASCII punct char at a token start: an alt-1 absorb may reach into + /// this batch. + alt1: u64, + /// P1 is a CJK-range char. Its all-zero carries say "closed unit", which is right for the + /// char AFTER it; bytes of the char itself straddling into the batch stay unclassified and + /// must defer (see `batch_masks`). + cjk: bool, +} + +fn carries(text: &[u8], tags: &[u8], scan: usize) -> Carries { + let mut cr = Carries::default(); + if scan == 0 { + return cr; + } + let p1 = char_lead(tags, scan - 1); + if (0xE3..=0xE9).contains(&text[p1]) { + // A CJK-range char is a closed unit: everything after it starts fresh, which is what + // all-zero carries say. + cr.cjk = true; + return cr; + } + let t1 = tags[p1]; + match t1 & 0x0F { + LET | MRK if is_member(t1) => cr.pl = 1, + NW | NO => cr.pd = 1, + SPC => { + cr.pws = 1; + cr.ps = 1; + } + WSO => { + cr.pws = 1; + cr.pwt = 1; + } + NLN => cr.pws = 1, + NMO | CTL => cr.pgap = 1, + _ if t1 == ZWJ => cr.pgap = 1, + _ => cr.po = 1, + } + if text[scan - 1].is_ascii_punctuation() { + // Was P1 (one byte) at a token start? Not when it continues a punct run or follows a + // space; a CJK-range P2 is a closed unit, so P1 starts fresh after it. + let at_start = if p1 == 0 { + true + } else { + let p2 = char_lead(tags, p1 - 1); + let t2 = tags[p2]; + (0xE3..=0xE9).contains(&text[p2]) + || !(t2 & 0x0F == SPC || in_mask(t2, mask::PUNCT_SYM) || t2 == ASM) + }; + cr.alt1 = u64::from(at_start); + } + cr +} + +fn batch_masks(text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + debug_assert!(scan + 64 < tags.len() && tags.len() == text.len()); + // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the blocks' load contract. + let tb = unsafe { Block::load(text, scan) }; + // CJK-range chars (lead 0xE3..=0xE9, a superset of the Split-2 ranges) are closed units + // that every other rule stops at; a batch containing any defers whole, BEFORE the tag + // masks are built. On CJK-dominated text the scanner degenerates to the scalar scan plus + // this one test (see the module doc). + if tb.range_full(0xE3, 6) != 0 { + return (0, u64::MAX); + } + // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract. + let blk = unsafe { Block::load(tags, scan) }; + if blk.any_range_tag(13, 1) { + return (0, u64::MAX); + } + let l0 = blk.eq_tag(LET); + let mk0 = blk.eq_full(MRK); // true marks join deepseek letter runs + let o0 = blk.range_tag(CON, 3) | blk.eq_full(ASM); // `[\p{P}\p{S}]` = Con|Pun|Apo|Sym (+ASM) + let gap0 = blk.range_tag(NMO, 1) | blk.eq_full(ZWJ); + let d0 = blk.range_tag(NW, 1); + let nl = blk.eq_tag(NLN); + let s = blk.eq_tag(SPC); + let wt0 = blk.eq_tag(WSO); + let c = if blk.any_eq_full(CONT) { + blk.eq_full(CONT) + } else { + 0 + }; + let pa = tb.ascii_punct(); + let alpha = tb.ascii_alpha(); + + let cr = carries(text, tags, scan); + + let lml0 = l0 | mk0; + let (mut lm, mut d, mut wt, mut g) = (lml0, d0, wt0, gap0); + let (c2, c3) = cont_runs(c); + if c != 0 { + if c & 1 != 0 { + let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1); + lm |= lead_in * cr.pl; + d |= lead_in * cr.pd; + wt |= lead_in * cr.pwt; + g |= lead_in * cr.pgap; + } + lm = fill(lm, c, c2, c3); + d = fill(d, c, c2, c3); + wt = fill(wt, c, c2, c3); + g = fill(g, c, c2, c3); + } + let lead = !c; + let len1 = lead & !(c >> 1); + let len2 = lead & (c >> 1) & !(c >> 2); + let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3); + let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3); + let w2 = wt0 & len2; + let w3 = wt0 & len3; + + let ws_f = s | nl | wt; + // The complement is the punct-run class here too (gap and CJK bytes land in it; both are + // covered by their own masks or the bad cover below). + let o = !(lm | d | ws_f | g); + + let mut bad = 0u64; + if cr.cjk && c & 1 != 0 { + // A CJK char straddling into the batch: its continuation bytes carry no class (the + // all-zero carries only cover the char after it), so they and the char right after + // them defer. + let e1 = (!c).trailing_zeros(); + bad |= (c & ((1u64 << e1) - 1)) | (1u64 << e1); + } + + // --- Letters: `[^\r\n\p{L}\p{P}\p{S}]?[\p{L}\p{M}]+` and alt-1 `[ascii-punct][A-Za-z]+` --- + let p_lm = (lm << 1) | cr.pl; + let p_s = (s << 1) | cr.ps; + let p_wt = (wt << 1) | cr.pwt; + let p_g = (g << 1) | cr.pgap; + let p_o = (o << 1) | cr.po; + + // --- Punct: ` ?[\p{P}\p{S}]+[\r\n]*` ------------------------------------------------------- + let b_punct = o0 & !p_o & !p_s; + let abs_seed = nl & ((o << 1) | cr.po); + let abs_n = if abs_seed == 0 { + 0 + } else { + smear_up(abs_seed, nl) + }; + let ws_eff = ws_f & !abs_n; + + // alt-1 absorbs: an ASCII-letter byte right after a token-starting ASCII punct char, and + // the rest of that `[A-Za-z]+` run. Where the run's end meets a letter-run member the two + // letter rules diverge: defer that bit. A run touching the batch end defers too (the next + // batch cannot tell an alt-1 run from a plain letter run). + let absorb_a = alpha & (((pa & b_punct) << 1) | cr.alt1); + let b_letters = lml0 & !p_lm & !p_s & !p_wt & !p_g & !absorb_a; + if absorb_a != 0 { + let zone = smear_up(absorb_a, alpha); + bad |= (zone << 1) & !alpha & lml0; + if zone >> 63 != 0 { + bad |= 1u64 << 63; + } + } + + // --- Gap runs: boundary at the run start, and at the last char before an absorbed letter + // run (the prefix split; one and the same bit for a single-char gap). The prefix split + // reads the NEXT char's lead, so a gap char whose follower sits past the batch edge + // defers. + let b_gap = (gap0 & !p_g) + | (gap0 + & ((len1 & (lml0 >> 1)) + | (len2 & (lml0 >> 2)) + | (len3 & (lml0 >> 3)) + | (len4 & (lml0 >> 4)))); + bad |= (gap0 & len1 & (u64::MAX << 63)) + | (gap0 & len2 & (u64::MAX << 62)) + | (gap0 & len3 & (u64::MAX << 61)) + | (gap0 & len4 & (u64::MAX << 60)); + + // --- Digits: `\p{N}{1,3}` (the cl100k cap-3 machinery) ------------------------------------- + let d_ascii = d0 & len1; + let dmb = (d & c) | (d0 & !len1); + let b_digits = if d_ascii & (d_ascii >> 1) != 0 { + digit_run_splits3(d_ascii) + } else { + d_ascii + }; + bad |= dmb | dmb << 1 | dmb >> 1; + + // --- Whitespace ----------------------------------------------------------------------------- + // The give-back is gated on the follower: a digit or CJK char after the run means the whole + // run is one token (`iso` covers the follower's bytes; CJK is inside `bad` anyway, but the + // gate keeps the algebra honest about why). + let la = tags[scan + 64] & 0x0F; + let la_iso = matches!(la, NW | NO) || (0xE3..=0xE9).contains(&text[scan + 64]); + let nn64 = la != CONT && !in_mask(la, mask::WS); + let nn64m = u64::from(nn64 && !la_iso).wrapping_neg(); + if abs_n >> 63 != 0 && !nn64 { + bad |= 1u64 << 63; + } + let nonws = !ws_eff; + if ws_eff >> 63 != 0 && !nn64 { + if nonws == 0 { + return (0, u64::MAX); + } + let h = 63 - nonws.leading_zeros(); + bad |= u64::MAX << (h + 1); + } + let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd); + if seed != 0 { + bad |= smear_up(seed, d_ascii); + } + + // A whitespace run followed by a digit keeps its last char (no give-back); the CJK case + // never reaches this path (any in-batch CJK deferred above), leaving only the lookahead. + let iso = d; + let nonws_ni = nonws & !iso; + let p_ws = (ws_eff << 1) | cr.pws; + let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff; + let ws_leads = (ws_leads1 | w2 | w3) & !abs_n; + let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61)); + let split_ok = (ws_leads1 & (nonws_ni >> 1)) + | (w2 & (nonws_ni >> 2)) + | (w3 & (nonws_ni >> 3)) + | (edge_last & nn64m); + let mut b_ws = ws_leads & (!p_ws | split_ok); + + let mut runs_n = nl & ws_eff & !bad; + while runs_n != 0 { + let f = runs_n.trailing_zeros(); + let below_gap = nonws & ((1u64 << f) - 1); + let a = if below_gap == 0 { + 0 + } else { + 64 - below_gap.leading_zeros() + }; + let e = (nonws & (u64::MAX << f)).trailing_zeros(); + let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e); + b_ws &= !run_mask; + b_ws |= 1u64 << a; + let q = 63 - (nl & run_mask).leading_zeros(); + // The post-newline tail: no give-back when the run's follower is a digit or CJK char + // (the whole tail is then one token). + let follower_iso = if e >= 64 { la_iso } else { iso >> e & 1 != 0 }; + if q + 1 < e { + b_ws |= 1u64 << (q + 1); + if !follower_iso { + let tail = run_mask & (u64::MAX << (q + 1)); + let tail_leads = ws_leads & tail; + b_ws |= 1u64 << (63 - tail_leads.leading_zeros()); + } + } + runs_n &= !run_mask; + } + + let boundary = b_letters | b_digits | b_punct | b_gap | b_ws; + (boundary & !bad, bad) +} diff --git a/tokenizers/atomsplit/src/fsm/masked/o200k.rs b/tokenizers/atomsplit/src/fsm/masked/o200k.rs new file mode 100644 index 0000000000..dbc77c4775 --- /dev/null +++ b/tokenizers/atomsplit/src/fsm/masked/o200k.rs @@ -0,0 +1,473 @@ +//! Masked scheme for the o200k regex family (o200k / GPT-4o / gpt-oss with contractions and +//! digit cap 3, Mistral tekken without contractions at cap 1) — the boundary algebra via +//! gigatoken's `o200k_family.rs` (MIT) over tag-fed class masks. +//! +//! Differences from the cl100k family: +//! - Letter runs are case-structured. Under leftmost-greedy backtracking the two letter +//! alternatives reduce to a phase automaton: a strict-upper char (`\p{Lu}\p{Lt}`) ends a +//! token exactly when the previous char is strict-lower (`\p{Ll}`) — "camelCase" splits +//! `camel|Case`, "HTTPResponse" stays one token. A strict-upper after a CASELESS letter +//! needs the phase and lookahead (the backtrack to the last caseless char), so those chars +//! defer to the scalar path. +//! - Contractions are attached suffixes of letter tokens, not a standalone alternative: +//! "don't" is ONE token and the char after a consumed suffix always starts a new one +//! ("can'ts" is `can't|s`). A contraction applies only when the apostrophe directly follows +//! a letter-run char; elsewhere `'` is ordinary punctuation. +//! - Punct runs absorb a `[\r\n/]*` tail. `/` is itself punct, so an absorbed tail always +//! begins with a newline; whether a batch-leading `[\r\n/]` run continues such a tail is +//! resolved by a bounded walkback over the preceding text. +//! - Marks (`\p{M}`) are dual-class: they join letter runs AND continue punct runs, so their +//! effective class is run-contextual. Mark chars (rare) defer to the scalar path with a bad +//! smear wide enough (8 bytes forward) to cover every boundary their class can influence +//! (two chars of multi-byte followers). + +use super::super::o200k::advance_o200k; +use super::block::Block; +use super::{MaskedFsm, char_lead, cont_runs, digit_run_splits3, fill, shl_sat, smear_up}; +use crate::fsm::{APO, ASM, Atom, CONT, LET, MRK, NLN, NO, NW, SPC, WSO, ZWJ, in_mask, mask}; + +pub(super) struct O200kMasked; + +impl MaskedFsm + for O200kMasked +{ + #[inline(always)] + fn batch_masks(&self, text: &[u8], tags: &[u8], scan: usize) -> (u64, u64) { + batch_masks::(text, tags, scan) + } + + #[inline(always)] + fn advance(&self, text: &[u8], tags: &[u8], i: usize, end: usize) -> usize { + advance_o200k::(text, tags, i, end) + } +} + +#[inline(always)] +fn is_tail_byte(b: u8) -> bool { + matches!(b, b'\r' | b'\n' | b'/') +} + +#[inline(always)] +fn is_member_mark(t: u8) -> bool { + t & 0x0F == MRK && t != ASM && t != ZWJ +} + +/// Was the tail-class byte at `scan - 1` absorbed by a punct run's `[\r\n/]*` tail (as opposed +/// to being a fresh punct-run `/` or a ws-run newline)? Walks the tail-class run back (bounded) +/// and classifies the char before it. `None`: unresolved (over-long run, or a preceding mark +/// whose own class is run-contextual). +fn prev_tail_absorbed(text: &[u8], tags: &[u8], scan: usize) -> Option { + debug_assert!(scan >= 1 && is_tail_byte(text[scan - 1])); + let mut r = scan - 1; + let mut steps = 0; + while r > 0 && is_tail_byte(text[r - 1]) { + r -= 1; + steps += 1; + if steps > 8 { + return None; + } + } + // T-run = text[r..scan]. The `[\r\n/]*` tail is greedy, so once absorption triggers — at + // the first newline that directly follows a punct-run char (an in-run slash, or the + // pre-run char for a run-leading newline) — everything to the run's end is absorbed. + // Before the trigger, newlines are ws-run members and slashes ordinary punct-run bytes. + let run = &text[r..scan]; + let mut trigger = usize::MAX; + let mut seen_slash = false; + for (j, &b) in run.iter().enumerate() { + if b == b'/' { + seen_slash = true; + continue; + } + if seen_slash { + trigger = j; + break; + } + if j == 0 { + if r == 0 { + continue; + } + let t = tags[char_lead(tags, r - 1)]; + if is_member_mark(t) || t & 0x0F >= 13 { + // A mark continues whatever run precedes it; Sentinel/MultiByte are opaque. + return None; + } + if !in_mask(t, mask::LETTER | mask::NUMBER | mask::WS) { + trigger = 0; + break; + } + } + } + Some(scan - 1 - r >= trigger) +} + +/// Two-back "punct or space" test for the char whose lead is `p2`. A slash may be an absorbed +/// tail byte (a token end, neither punct-run member nor space), so it resolves through the +/// walkback. `None`: unresolved (the caller sets `force_bad_lead`). A mark P2 answers 0: the +/// bits that read it wrongly sit inside the previous batch's mark smear, and the scalar +/// overrun from there covers them (the walker's resume masking). +fn c2_os_at(text: &[u8], tags: &[u8], p2: usize) -> Option { + if text[p2] == b'/' { + return prev_tail_absorbed(text, tags, p2 + 1).map(|abs| u64::from(!abs)); + } + let t2 = tags[p2]; + Some(u64::from( + t2 & 0x0F == SPC || (!is_member_mark(t2) && in_mask(t2, mask::NOT_WS_L_N)), + )) +} + +/// Boundary carries from the two chars before the batch (the cl100k set, plus the case classes +/// and the absorbed-tail resolution). +#[derive(Default)] +struct Carries { + pl: u64, + pu: u64, + pcl: u64, + ps: u64, + pwt: u64, + po: u64, + pws: u64, + pd: u64, + /// P1 is a member-mark (seeds the mark smear at bit 0). + pmk: u64, + c2_os: u64, + b2b_in: u64, + /// P1 is an absorbed `[\r\n/]*` tail byte whose token may continue into this batch. + p_abs: bool, + /// The tail walkback could not resolve: the batch's leading tail-class run (plus the byte + /// after it) can't be trusted. + force_bad_lead: bool, +} + +fn carries(text: &[u8], tags: &[u8], scan: usize, c: u64) -> Carries { + let mut cr = Carries::default(); + if scan == 0 { + return cr; + } + if is_tail_byte(text[scan - 1]) { + // An absorbed tail ended the previous token, so every "P1 is X" carry is zero and only + // the tail-continuation seed survives. A fresh `/` is an ordinary punct byte; fresh + // `\r\n` are ws-run newlines. + match prev_tail_absorbed(text, tags, scan) { + None => cr.force_bad_lead = true, + Some(true) => cr.p_abs = true, + Some(false) => { + if text[scan - 1] == b'/' { + cr.po = 1; + } else { + cr.pws = 1; + } + match c2_os_at(text, tags, char_lead(tags, scan - 2)) { + Some(v) => cr.c2_os = v, + None => cr.force_bad_lead = true, + } + } + } + return cr; + } + let p1 = char_lead(tags, scan - 1); + let c2v = if p1 == 0 { + Some(0) + } else { + c2_os_at(text, tags, char_lead(tags, p1 - 1)) + }; + match c2v { + Some(v) => { + if c & 1 != 0 { + cr.b2b_in = v << (!c).trailing_zeros(); + } else { + cr.c2_os = v; + } + } + None => cr.force_bad_lead = true, + } + let t1 = tags[p1]; + match t1 & 0x0F { + LET => { + cr.pl = 1; + cr.pu = u64::from(t1 == Atom::UpperLetter as u8); + cr.pcl = u64::from(t1 == Atom::Letter as u8); + } + NW | NO => cr.pd = 1, + SPC => { + cr.pws = 1; + cr.ps = 1; + } + WSO => { + cr.pws = 1; + cr.pwt = 1; + } + MRK if is_member_mark(t1) => cr.pmk = 1, + _ => cr.po = 1, + } + cr +} + +fn batch_masks( + text: &[u8], + tags: &[u8], + scan: usize, +) -> (u64, u64) { + debug_assert!(scan + 64 < tags.len() && tags.len() == text.len()); + // SAFETY: `scan + 64 < tags.len()` (walker guarantee), the block's load contract. + let blk = unsafe { Block::load(tags, scan) }; + if blk.any_range_tag(13, 1) { + return (0, u64::MAX); + } + let l0 = blk.eq_tag(LET); + let ub0 = blk.eq_full(Atom::UpperLetter as u8); + let cl0 = blk.eq_full(Atom::Letter as u8); + let mk0 = blk.eq_full(MRK); // true marks: refined Mark tags (ASM/ZWJ) are punct-class + let d0 = blk.range_tag(NW, 1); + let nl = blk.eq_tag(NLN); + let s = blk.eq_tag(SPC); + let wt0 = blk.eq_tag(WSO); + let ap = if blk.any_eq_full(APO) { + blk.eq_full(APO) + } else { + 0 + }; + let c = if blk.any_eq_full(CONT) { + blk.eq_full(CONT) + } else { + 0 + }; + + let cr = carries(text, tags, scan, c); + + let (mut l, mut u, mut clb, mut d, mut wt, mut mk) = (l0, ub0, cl0, d0, wt0, mk0); + let (c2, c3) = cont_runs(c); + if c != 0 { + if c & 1 != 0 { + let lead_in = c & ((1u64 << (!c).trailing_zeros()) - 1); + l |= lead_in * cr.pl; + u |= lead_in * cr.pu; + clb |= lead_in * cr.pcl; + d |= lead_in * cr.pd; + wt |= lead_in * cr.pwt; + mk |= lead_in * cr.pmk; + } + l = fill(l, c, c2, c3); + u = fill(u, c, c2, c3); + clb = fill(clb, c, c2, c3); + d = fill(d, c, c2, c3); + wt = fill(wt, c, c2, c3); + mk = fill(mk, c, c2, c3); + } + mk |= cr.pmk; // a mark P1 poisons bit 0 even when it ends exactly at the batch edge + let lead = !c; + let len1 = lead & !(c >> 1); + let len2 = lead & (c >> 1) & !(c >> 2); + let len3 = lead & (c >> 1) & (c >> 2) & !(c >> 3); + let len4 = lead & (c >> 1) & (c >> 2) & (c >> 3); + let w2 = wt0 & len2; + let w3 = wt0 & len3; + + let ws_f = s | nl | wt; + // Marks land in the complement (dual-class); every bit that can read them is inside the + // mark bad smear below, so their punct-class reading is never trusted. + let o = !(l | d | ws_f); + + // --- Absorbed `[\r\n/]*` tails --------------------------------------------------------- + // The tail class needs the slash mask from the TEXT block; skip that load when the batch + // has no newline and no tail context carried in. + let (tcls, abs_t) = if nl != 0 || cr.p_abs || cr.force_bad_lead { + // SAFETY: `scan + 64 < text.len()` (walker guarantee), the block's load contract. + let sl = unsafe { Block::load(text, scan) }.eq_full(b'/'); + let tcls = nl | sl; + let abs_seed = (nl & ((o << 1) | cr.po)) | (u64::from(cr.p_abs) & tcls); + let abs_t = if abs_seed == 0 { + 0 + } else { + smear_up(abs_seed, tcls) + }; + (tcls, abs_t) + } else { + (nl, 0) + }; + let ob_eff = o & !abs_t; + + // --- Letters (see the cl100k scheme for the base rules) --------------------------------- + let c_test = ((ob_eff | s) << 1) | cr.po | cr.ps; + let b2back = ((c_test & len1) << 1) + | ((c_test & len2) << 2) + | ((c_test & len3) << 3) + | ((c_test & len4) << 4) + | cr.c2_os + | cr.b2b_in; + let p_l = (l << 1) | cr.pl; + let p_u = (u << 1) | cr.pu; + let p_cl = (clb << 1) | cr.pcl; + let p_s = (s << 1) | cr.ps; + let p_wt = (wt << 1) | cr.pwt; + let p_o = (ob_eff << 1) | cr.po; + let absorb = p_o & !b2back; + // Casing boundary: a strict-upper char after a strict-lower one. (For ASCII text this is + // the whole rule; upper-after-caseless defers below.) + let p_sl = p_l & !p_u & !p_cl; + let b_letters = (l0 & !p_l & !p_s & !p_wt & !absorb) | (ub0 & p_sl); + + // --- Digits ------------------------------------------------------------------------------ + let d_ascii = d0 & len1; + let dmb = (d & c) | (d0 & !len1); + let b_digits = if DIGIT_CAP == 3 { + if d_ascii & (d_ascii >> 1) != 0 { + digit_run_splits3(d_ascii) + } else { + d_ascii + } + } else { + d0 // cap 1: every digit char its own token + }; + + // --- Punct: ` ?[^\s\p{L}\p{N}]+[\r\n/]*` -------------------------------------------------- + let b_punct = (ob_eff & lead) & !p_o & !p_s; + + // --- Bad zones ---------------------------------------------------------------------------- + let mut bad = if DIGIT_CAP == 3 { + dmb | dmb << 1 | dmb >> 1 + } else { + 0 + }; + if mk != 0 { + // A mark's run-contextual class can affect boundaries up to two chars after it (8 + // bytes of multi-byte followers) and the byte before. + bad |= mk + | (mk << 1) + | (mk << 2) + | (mk << 3) + | (mk << 4) + | (mk << 5) + | (mk << 6) + | (mk << 7) + | (mk << 8) + | (mk >> 1); + } + // A strict-upper char after a caseless letter: phase- and lookahead-dependent. + bad |= ub0 & ((clb << 1) | cr.pcl); + if cr.force_bad_lead { + bad |= (smear_up(tcls & 1, tcls) << 1) | 0b11; + } + + // --- Whitespace --------------------------------------------------------------------------- + let ws_eff = ws_f & !abs_t; + let la = tags[scan + 64] & 0x0F; + let nn64 = la != CONT && !in_mask(la, mask::WS); + let nn64m = u64::from(nn64).wrapping_neg(); + + // An absorbed tail touching the batch end continues iff byte 64 is tail-class; the next + // batch's tail walkback re-derives the context either way, so nothing defers here. A ws + // run touching the batch end still defers when byte 64 is ws. + let nonws = !ws_eff; + if ws_eff >> 63 != 0 && !nn64 { + if nonws == 0 { + return (0, u64::MAX); + } + let h = 63 - nonws.leading_zeros(); + bad |= u64::MAX << (h + 1); + } + if DIGIT_CAP == 3 { + let seed = (d_ascii & (bad << 1)) | (d_ascii & cr.pd); + if seed != 0 { + bad |= smear_up(seed, d_ascii); + } + } + + let ws_leads1 = (s | nl | (wt0 & len1)) & ws_eff; + let ws_leads = (ws_leads1 | w2 | w3) & !abs_t; + let p_ws = (ws_eff << 1) | cr.pws; + let edge_last = (ws_leads1 & (1 << 63)) | (w2 & (1 << 62)) | (w3 & (1 << 61)); + let split_ok = (ws_leads1 & (nonws >> 1)) + | (w2 & (nonws >> 2)) + | (w3 & (nonws >> 3)) + | (edge_last & nn64m); + let mut b_ws = ws_leads & (!p_ws | split_ok); + + let mut runs_n = nl & ws_eff & !bad; + while runs_n != 0 { + let f = runs_n.trailing_zeros(); + let below_gap = nonws & ((1u64 << f) - 1); + let a = if below_gap == 0 { + 0 + } else { + 64 - below_gap.leading_zeros() + }; + let e = (nonws & (u64::MAX << f)).trailing_zeros(); + let run_mask = (u64::MAX << a) & !shl_sat(u64::MAX, e); + b_ws &= !run_mask; + b_ws |= 1u64 << a; + let q = 63 - (nl & run_mask).leading_zeros(); + if q + 1 < e { + b_ws |= 1u64 << (q + 1); + let tail = run_mask & (u64::MAX << (q + 1)); + let tail_leads = ws_leads & tail; + b_ws |= 1u64 << (63 - tail_leads.leading_zeros()); + } + runs_n &= !run_mask; + } + + let mut boundary = b_letters | b_digits | b_punct | b_ws; + + // --- Contractions: suffix `(?i:'s|'t|'re|'ve|'m|'ll|'d)?` --------------------------------- + // An apostrophe at a boundary right after a letter-run char merges the suffix into that + // token and forces a boundary right after it. + if CONTRACTION { + let mut cand = ap & boundary & p_l & !bad; + let mut last_forced = usize::MAX; + while cand != 0 { + let i = cand.trailing_zeros() as usize; + cand &= cand - 1; + if i <= 2 { + // The preceding letter could itself end an earlier contraction that started + // before the batch: scalar. + bad |= 0b111u64 << i; + continue; + } + if i >= 61 { + bad |= u64::MAX << i; + break; + } + if i == last_forced { + // "x'll'd": the letter before this apostrophe is a consumed suffix's last + // char; a new (prefix) match starts here instead. + continue; + } + // The letter before this apostrophe may itself be a consumed suffix's last char + // resolved where `last_forced` can't see it (a scalar-walked zone, or a fixup + // before the batch): locally ambiguous, defer. + let p = scan + i; + let prev_suffix_possible = (text[p - 2] == b'\'' + && matches!(text[p - 1] | 0x20, b's' | b'd' | b'm' | b't')) + || (text[p - 3] == b'\'' + && (matches!( + (text[p - 2] | 0x20, text[p - 1] | 0x20), + (b'l', b'l') | (b'v', b'e') | (b'r', b'e') + ) || (text[p - 2] == 0xC5 && text[p - 1] == 0xBF))); + if prev_suffix_possible { + bad |= 0b111u64 << i; + continue; + } + let b1 = text[p + 1]; + if b1 >= 0x80 { + // `(?i:'s)` also matches 'ſ (U+017F): defer. + bad |= 0b111u64 << i; + continue; + } + let k = match b1 | 0x20 { + b's' | b'd' | b'm' | b't' => 2, + b'l' if text[p + 2] | 0x20 == b'l' => 3, + b'v' if text[p + 2] | 0x20 == b'e' => 3, + b'r' if text[p + 2] | 0x20 == b'e' => 3, + _ => 0, + }; + if k != 0 { + boundary &= !(1u64 << i); + boundary &= !(((1u64 << (k - 1)) - 1) << (i + 1)); + boundary |= 1u64 << (i + k); + last_forced = i + k; + } + } + } + + (boundary & !bad, bad) +} diff --git a/tokenizers/atomsplit/src/fsm/o200k.rs b/tokenizers/atomsplit/src/fsm/o200k.rs index de02fbd79f..6914e02c75 100644 --- a/tokenizers/atomsplit/src/fsm/o200k.rs +++ b/tokenizers/atomsplit/src/fsm/o200k.rs @@ -125,6 +125,143 @@ pub fn scan_tekken(text: &[u8], tags: &[u8], emit: impl FnMut(Span)) { o200k::(text, tags, emit); } +/// Is tag `t` a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in; coarse +/// `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w` +/// but not `[\p{L}\p{M}]`. +#[inline(always)] +fn member(t: u8) -> bool { + let c = t & 0x0F; + c == LET || (c == MRK && t != ASM && t != ZWJ) +} + +#[inline(always)] +fn is_lm(tags: &[u8], a: usize, end: usize) -> bool { + a < end && member(tags[a]) +} + +/// Maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`). +#[inline(always)] +fn letter_end(tags: &[u8], a: usize, end: usize) -> usize { + let mut p = a; + // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter` + // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run; only a + // coarse `Mark` lane pays the refinement test. Byte-exact with the plain scan below. + // SAFETY: `p + 16 <= end <= tags.len()` in the body. + while p + 16 <= end { + let mut brk = 16; + for k in 0..16 { + let t = unsafe { *tags.get_unchecked(p + k) }; + if t == CONT || t & 0x0F == LET { + continue; + } + if t & 0x0F == MRK && t != ASM && t != ZWJ { + continue; + } + brk = k; + break; + } + if brk < 16 { + return p + brk; + } + p += 16; + } + while p < end && (tags[p] == CONT || member(tags[p])) { + p += 1; + } + p +} + +/// rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if +/// none. `/` is in the `+` body too — the trailing class only matters after the `+` stops at a +/// `\r\n`. +#[inline(always)] +fn other(text: &[u8], tags: &[u8], sp0: usize, end: usize) -> usize { + let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); + if p > sp0 { + while p < end && (tags[p] == NLN || text[p] == b'/') { + p += char_len(text[p]); + } + } + p +} + +/// End of the token starting at `i` (`i < end`, `i` on a token boundary): one rule dispatch of +/// the o200k-family regex. A letter-path token is the FIRST case sub-run from its start (the +/// case split restarts at every sub-token, so one dispatch per sub-token equals the run loop in +/// [`emit_o200k_letters`]); the masked scanner re-derives tokens with this where its batch +/// masks are not trustworthy. +#[inline(always)] +pub(super) fn advance_o200k( + text: &[u8], + tags: &[u8], + i: usize, + end: usize, +) -> usize { + let one_letter = |ls: usize| -> usize { + let re = letter_end(tags, ls, end); + let e = o200k_letter_match(tags, ls, re); + if CONTRACTION && e == re { + e + contraction(text, e) + } else { + e + } + }; + let b = text[i]; + match tags[i] & 0x0F { + NW | NO => { + let (mut p, mut cnt) = (i, 0); + while p < end && cnt < DIGIT_CAP && in_mask(tags[p], mask::NUMBER) { + p += char_len(text[p]); + cnt += 1; + } + p + } + LET | MRK => { + if tags[i] != ASM && tags[i] != ZWJ { + one_letter(i) + } else { + let a = i + char_len(b); + if is_lm(tags, a, end) { + one_letter(a) + } else { + other(text, tags, i, end) + } + } + } + SPC => { + let a = i + 1; // Space is ASCII (0x20) + if is_lm(tags, a, end) { + one_letter(a) + } else { + let p = other(text, tags, a, end); + if p > a { + p + } else { + ws_tail(text, tags, i, end) + } + } + } + WSO => { + let a = i + char_len(b); + if is_lm(tags, a, end) { + one_letter(a) + } else { + ws_tail(text, tags, i, end) + } + } + NLN => ws_tail(text, tags, i, end), + CON | PUN | APO | SYM | NMO | CTL => { + let a = i + char_len(b); + if is_lm(tags, a, end) { + one_letter(a) + } else { + other(text, tags, i, end) + } + } + _ => i + char_len(b), + } +} + fn o200k( text: &[u8], tags: &[u8], @@ -136,61 +273,13 @@ fn o200k( // `tags[i]` in this fsm + its `run_end`/`letter_*` scans. (Callers guarantee `tags.len() >= end`.) let tags = &tags[..end]; - // Is tag `t` (at byte `p`) a real `[\p{L}\p{M}]` member? Coarse `Letter` (any case) is always in; - // coarse `Mark` is in only as a true `\p{M}` — ALPHA_SYM (`\p{S}`) and ZWJ/ZWNJ (`\p{Cf}`) are `\w` - // but not `[\p{L}\p{M}]`. The `ds_is_zwj` byte-peek is thus paid ONLY for Marks, never for letters. - let member = |t: u8, p: usize| -> bool { - let c = t & 0x0F; - let _ = p; - c == LET || (c == MRK && t != ASM && t != ZWJ) - }; - let is_lm = |a: usize| a < end && member(tags[a], a); - // maximal `[\p{L}\p{M}]+` run from `a` (byte-wise; continuation bytes ride along — see `run_end`). - let letter_end = |a: usize| -> usize { - let mut p = a; - // logos-style fast loop: 16 tags/chunk, one bounds check, unchecked reads. A plain `Letter` - // (low nibble 0, incl Han — o200k keeps all letters) or a `Cont` byte stays in-run with no - // `ds_is_zwj` peek; only a coarse `Mark` lane pays the peek. Byte-exact with the scalar scan. - // SAFETY: `p + 16 <= end <= tags.len()`/`text.len()` in the body. - while p + 16 <= end { - let mut brk = 16; - for k in 0..16 { - let t = unsafe { *tags.get_unchecked(p + k) }; - if t == CONT || t & 0x0F == LET { - continue; - } - if t & 0x0F == MRK && t != ASM && t != ZWJ { - continue; - } - brk = k; - break; - } - if brk < 16 { - return p + brk; - } - p += 16; - } - while p < end && (tags[p] == CONT || member(tags[p], p)) { - p += 1; - } - p - }; - // rule 4 `[^\s\p{L}\p{N}]+[\r\n/]*` from `sp0` (any leading space already consumed); `sp0` if none. - // `/` is in the `+` body too — the trailing class only matters after the `+` stops at a `\r\n`. - let other = |sp0: usize| -> usize { - let mut p = run_end(tags, sp0, end, mask::NOT_WS_L_N); - if p > sp0 { - while p < end && (tags[p] == NLN || text[p] == b'/') { - p += char_len(text[p]); - } - } - p - }; + let is_lm = |a: usize| is_lm(tags, a, end); // rules 5-7 (`\s*[\r\n]+ | \s+(?!\S) | \s+`) → the shared `ws_tail` (identical to cl100k). let ws = |i: usize| -> usize { ws_tail(text, tags, i, end) }; + let other = |sp0: usize| -> usize { other(text, tags, sp0, end) }; // The letter rules: case-split the run starting at `ls`, first sub-token starting at the prefix `pfx`. let letters = |pfx: usize, ls: usize, emit: &mut E| -> usize { - emit_o200k_letters::(text, tags, pfx, ls, letter_end(ls), emit) + emit_o200k_letters::(text, tags, pfx, ls, letter_end(tags, ls, end), emit) }; let mut i = 0; diff --git a/tokenizers/atomsplit/tests/fsm.rs b/tokenizers/atomsplit/tests/fsm.rs index 67e0c92627..3aa75ab7fc 100644 --- a/tokenizers/atomsplit/tests/fsm.rs +++ b/tokenizers/atomsplit/tests/fsm.rs @@ -2,9 +2,47 @@ use atomsplit::classify::{classify, mask}; use atomsplit::fsm::{ Span, class_runs_into, emit_class_spans, fsm_byte_level, fsm_cl100k, fsm_deepseek, fsm_o200k, - fsm_tekken, + fsm_tekken, scan_byte_level, scan_byte_level_masked, scan_cl100k_cap, scan_cl100k_cap_masked, + scan_deepseek, scan_deepseek_masked, scan_o200k, scan_o200k_masked, scan_tekken, + scan_tekken_masked, }; +/// The shared sweep for a masked/scalar scanner pair: compare spans on the corpus behind +/// every 64-byte-edge offset (leading padding 0..=70), and with truncations exercising the +/// scalar tail at every remaining length. +type ScanFn = dyn Fn(&[u8], &[u8], &mut dyn FnMut(Span)); + +fn masked_matches_scalar(corpus: &str, scalar: &ScanFn, masked: &ScanFn) { + fn spans_of(scan: &ScanFn, s: &[u8]) -> Vec { + let mut tags = vec![0u8; s.len()]; + classify(s, &mut tags); + let mut v = Vec::new(); + scan(s, &tags, &mut |sp| v.push(sp)); + v + } + let check = |s: &str| { + let a = spans_of(scalar, s.as_bytes()); + let b = spans_of(masked, s.as_bytes()); + assert_eq!(b, a, "input len {}: {:?}", s.len(), s); + }; + for pad in 0..=70 { + check(&format!("{}{}", "x".repeat(pad), corpus)); + } + for pad in [0, 37] { + let padded = format!("{}{}", "x".repeat(pad), corpus); + for len in padded.len().saturating_sub(140)..=padded.len() { + if padded.is_char_boundary(len) { + check(&padded[..len]); + } + } + } + for len in 0..=70 { + if corpus.is_char_boundary(len) { + check(&corpus[..len]); + } + } +} + /// Run a no-push fsm into a fresh buffer and return the emitted spans. fn spans(f: impl Fn(&[u8], &[u8], &mut [Span]) -> usize, s: &str) -> Vec { let mut tags = vec![0u8; s.len()]; @@ -48,6 +86,111 @@ fn deepseek_rules() { assert_eq!(ds("!!!"), vec![(0, 3)]); // \p{P}∪\p{S} run } +/// Byte-exactness gate for the masked byte-level scanner: `scan_byte_level_masked` must emit the +/// spans `scan_byte_level` emits, on every input. The corpus stresses every rule shape the batch +/// algebra rewrites: contractions in both cases and at rejection edges, prefix spaces before all +/// three run classes, unbounded numbers including non-ASCII digits, multi-byte whitespace (the +/// bad-zone route), CJK/emoji/ZWJ runs, CRLF and tab runs, and runs longer than a batch. The +/// padding sweep (0..=70 leading bytes) moves every shape across every 64-byte batch edge, so +/// batch-edge carries, the bit-63 lookahead and every bad-zone route are hit at every offset; the +/// truncation sweep exercises the scalar tail at every remaining length. +#[test] +fn masked_scan_matches_scan_byte_level() { + let corpus = concat!( + "I'm 12345 ok, don't they'll 've 'lx x's ''s IT'S 's mid'dle end' ", + "hello world\r\nmore\ttabs\t\t spaced end ", + "no.1 中文漢字テスト مرحبا १२३४५६७८९० ¹²³ ½¾ ", + "emoji 😀😀 zwj 👩\u{200d}🔬 nbsp\u{a0}x thin\u{2009}y wide\u{3000}z ", + "((()))!!!??? #$%&' “curly” apostrophe’d ", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ", + "9999999999999999999999999999999999999999999999999999999999999999999999999999999 ", + "中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中 end", + ); + masked_matches_scalar( + corpus, + &|t, tg, e| scan_byte_level(t, tg, e), + &|t, tg, e| scan_byte_level_masked(t, tg, e), + ); +} + +/// Byte-exactness gate for the masked cl100k-family scanner, at every shipped digit cap. On top +/// of the byte-level shapes: capped digit runs (pure ASCII, pure Devanagari, and mixed, so the +/// char-counted bad route is hit mid-run), letters after punct at run starts vs mid-run (the +/// two-chars-back absorb test), newlines absorbed after punct runs (`[\r\n]*`), whitespace runs +/// with interior newlines (`\s*[\r\n]`), tab prefixes, and an apostrophe before a non-ASCII +/// char (the `'ſ` defer). +#[test] +fn masked_scan_matches_scan_cl100k_cap() { + let corpus = concat!( + "I'm 12345 ok, don't they'LL 've 'lx x's ''s IT'S 's 3'ts end' ", + "a1234 12٣45 ١٢٣٤٥٦ १२३४५६७८९० 999999999999999999999999999999999999999999999999 ", + "!!a x!a ?!x. ;\n\n};\r\n\r\n foo();\nbar() //x\n\n\n end\n", + "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x ", + "no.1 中文漢字テスト مرحبا nbsp\u{a0}x thin\u{2009}y wide\u{3000}z 'ſok “curly”’d ", + "((()))!!!??? #$%&' 😀😀 👩\u{200d}🔬 ", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa ", + "中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中中 end", + ); + for cap in [3, 1, usize::MAX] { + masked_matches_scalar( + corpus, + &move |t, tg, e| scan_cl100k_cap(t, tg, cap, e), + &move |t, tg, e| scan_cl100k_cap_masked(t, tg, cap, e), + ); + } +} + +/// Byte-exactness gate for the masked o200k/tekken scanner. On top of the cl100k shapes: case +/// splits (camelCase, all-upper runs, upper after caseless — the deferred backtrack), suffix +/// contractions incl. chains ("can'ts", "x'll'd") and prefix apostrophes after digits, `[\r\n/]*` +/// tails with slash runs before and after newlines (the walkback shapes), and combining marks +/// (run-contextual class, the wide bad smear). +#[test] +fn masked_scan_matches_scan_o200k_and_tekken() { + let corpus = concat!( + "camelCase HTTPResponse XMLHttpRequest AAAA aaaa aA Aa 中B B中b ʰupper Xʰa 中中中中 ", + "don't they'LL CAN'TS x'll'd 3'ts I'm'll a'sit 's'' end' 'ſok ", + "a1234 12٣45 १२३४५६ 99999999999999999999999999999999999999999999999999999999999999999 ", + "foo();\nbar() .\n// ///x\r\n/// a/b//c\n\n//d e\u{301}f g\u{5bf}h zwj\u{200c}x ", + "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x nbsp\u{a0}x wide\u{3000}z ", + "((()))!!!??? #$%&' “curly”’d 😀😀 مرحبا מבחן ", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa end" + ); + masked_matches_scalar( + corpus, + &|t, tg, e| scan_o200k(t, tg, e), + &|t, tg, e| scan_o200k_masked(t, tg, e), + ); + masked_matches_scalar( + corpus, + &|t, tg, e| scan_tekken(t, tg, e), + &|t, tg, e| scan_tekken_masked(t, tg, e), + ); +} + +/// Byte-exactness gate for the masked deepseek scanner. On top of the shared shapes: CJK +/// letter/punct runs and their neighborhoods (the closed-unit rule and the bad cover), gap runs +/// (controls / NumericOther / ZWJ) with and without a following letter run (the prefix split), +/// alt-1 `[ascii-punct][A-Za-z]+` including collisions with non-ASCII letters ("_naïve"), and +/// whitespace runs followed by digits or CJK (no give-back). +#[test] +fn masked_scan_matches_scan_deepseek() { + let corpus = concat!( + "abc中def 中文漢字テスト!ひらがな・カタカナ 中中中中中中中 mixed中123中ok 拼音列表(e.g. 表!x ", + "_abc (foo) .py x!a _naïve _né !!a a/b.c 3'ts don't ½x ¼¼y\u{7f}z zwj\u{200c}gap\u{200c}\u{200c}word ", + "a1234 12٣45 999999999999999999999999999999999999999999999999999999999999999999 12 中 ", + "x 123 y\t\t\t45 z \n 99 w 中 Model):\n money = f(x):\n\n\t indent ", + "hello world\r\nmore\ttabs\t\t spaced \ta \r\n \n\t\r\n\t x wide\u{3000}z 中 12\n\n34 ", + "((()))!!!??? #$%&' “curly”’d 😀😀 مرحبا ", + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa end" + ); + masked_matches_scalar( + corpus, + &|t, tg, e| scan_deepseek(t, tg, e), + &|t, tg, e| scan_deepseek_masked(t, tg, e), + ); +} + #[test] fn byte_level_rules() { let bl = |s| spans(fsm_byte_level, s); diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index ec2d397e3c..92bd56de48 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -906,26 +906,27 @@ impl PipelineTokenizer { classify(bytes, &mut tags[..bytes.len()]); let tags = &tags[..bytes.len()]; use atomsplit::fsm::{ - scan_byte_level, scan_cl100k_cap, scan_deepseek, scan_o200k, scan_tekken, + scan_byte_level_masked, scan_cl100k_cap_masked, scan_deepseek_masked, + scan_o200k_masked, scan_tekken_masked, }; // One emit closure literal per arm: each scan gets its own instance by // value, which is what lets the emit inline into the scan loop. match scan { - FusedScan::Gpt(GptFsm::Gpt2) => scan_byte_level(bytes, tags, |span| { + FusedScan::Gpt(GptFsm::Gpt2) => scan_byte_level_masked(bytes, tags, |span| { model.tokenize_span(&chunk[span.range()], scratch, output); }), FusedScan::Gpt(GptFsm::Cl100k { digit_cap }) => { - scan_cl100k_cap(bytes, tags, digit_cap, |span| { + scan_cl100k_cap_masked(bytes, tags, digit_cap, |span| { model.tokenize_span(&chunk[span.range()], scratch, output); }) } - FusedScan::Gpt(GptFsm::O200k) => scan_o200k(bytes, tags, |span| { + FusedScan::Gpt(GptFsm::O200k) => scan_o200k_masked(bytes, tags, |span| { model.tokenize_span(&chunk[span.range()], scratch, output); }), - FusedScan::Gpt(GptFsm::Tekken) => scan_tekken(bytes, tags, |span| { + FusedScan::Gpt(GptFsm::Tekken) => scan_tekken_masked(bytes, tags, |span| { model.tokenize_span(&chunk[span.range()], scratch, output); }), - FusedScan::DeepSeek => scan_deepseek(bytes, tags, |span| { + FusedScan::DeepSeek => scan_deepseek_masked(bytes, tags, |span| { model.tokenize_span(&chunk[span.range()], scratch, output); }), } From ac507686a4fd7d893c7f3bc0b92f101770a351ff Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 03:08:28 +0200 Subject: [PATCH 12/13] perf(pipeline): read pending rewrites off the raw text instead of writing them A normalizer whose rewrite changes the text one character at a time does not have to write it: positions map one to one, so downstream stages can read the rewritten form off the raw text. `PendingRewrite` names that class of rewrite (every `from` reads as `to`, one may be prepended), and `Normalizer::pending_rewrite` lets every normalizer say whether it is one: the fused MetaspaceNormalizer (the SentencePiece space swap) and a single-character literal Replace are, everything else keeps writing. `NormalizedText::from_chain` runs the chain and hands a trailing pending rewrite back unwritten; `write` produces the text for consumers that need the bytes, and MetaspaceNormalizer's own swap now goes through the same writer. The consumer is the zero-copy encode path, built when a chain ends in the space swap over proven cuts into a char-atom BPE with no `normalized` added tokens (llama-2, gemma-4). `ZeroCopyMetaspace` finds the same cuts on the raw text (one memchr2 pass over `from` and `to`), and the model reads each raw span through a `CharSwap`, seeding the delimiter's id for every raw `from`. The word cache keys on the raw bytes; a rewritten form never holds a `from`, so overlapping keys always agree on their ids. Only the one word a prepend touches is still rewritten for real. The scan compares single bytes, not runtime-length slices: a slice compare of run-time width compiles to a memcmp call per candidate, measured at 15-20% of the whole win. The veto is taken at its prefilter's word: a cut whose preceding byte could open a veto piece is skipped instead of checked against rewritten bytes this path never builds. Skipping a cut never changes the ids, and the byte fires for under 1% of spaces on the worst corpus (gemma-4's one piece, code text). Same-binary interleaved A/B (examples/normalize_claims.rs, 512 kB per corpus, warm cache): llama-2 +10-64%, gemma-4 +29-75% encode throughput across eng/cmn/code at 10 kB and per-line inputs. Ids match the written path over every corpus, and the released-crate oracle stays green 9/9. Co-Authored-By: Claude Fable 5 --- .../tk-encode/examples/normalize_claims.rs | 337 ++++++++++++ tokenizers/tk-encode/src/models/bpe/model.rs | 78 ++- .../tk-encode/src/normalizers/metaspace.rs | 93 ++-- tokenizers/tk-encode/src/normalizers/mod.rs | 19 + .../tk-encode/src/normalizers/replace.rs | 18 + .../src/pre_tokenizers/proven_cuts.rs | 117 +++- .../tk-encode/src/tokenizer/pipeline.rs | 518 +++++++++++++++++- .../src/vocab/bucket_added_vocabulary.rs | 6 + 8 files changed, 1102 insertions(+), 84 deletions(-) create mode 100644 tokenizers/tk-encode/examples/normalize_claims.rs diff --git a/tokenizers/tk-encode/examples/normalize_claims.rs b/tokenizers/tk-encode/examples/normalize_claims.rs new file mode 100644 index 0000000000..75dd97cf9c --- /dev/null +++ b/tokenizers/tk-encode/examples/normalize_claims.rs @@ -0,0 +1,337 @@ +//! What the space rewrite costs when written, and what the zero-copy path wins back. +//! +//! For each SentencePiece-shaped model this runs the `encode_generic::` ablation +//! ladder (the same methodology as `fixture_bench`) over long and short inputs, then +//! decomposes the fused normalizer's pass into its three parts: counting the spaces, +//! allocating the rewrite `String`, and writing it. Models taking the zero-copy path +//! (see `ZeroCopyMetaspace`) get an interleaved A/B against the written rewrite, with +//! ids compared over the whole corpus first. A last probe measures what the second +//! special-token scan adds once a tokenizer holds `normalized` added tokens. + +use std::hint::black_box; +use std::path::Path; +use std::time::Instant; + +use atomsplit::literal::Literal; +use tk_encode::pipeline::{Model, PipelineTokenizer}; +use tk_encode::{AddedToken, NormalizerWrapper, Tokenizer}; + +const DATA: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../data"); +const REPS: usize = 9; +const CHUNK_BYTES: usize = 10 * 1024; +const MAX_BYTES: usize = 512 * 1024; + +/// (config file, prepend for the decomposition replica; `None` skips it because the +/// model runs the `drop_whitespace` path, which is a different rewrite) +const MODELS: &[(&str, Option)] = &[ + ("llama-2.json", Some(true)), + ("gemma-4.json", Some(false)), + ("t5-base.json", None), + ("albert-base-v1-tokenizer.json", None), +]; + +const FIXTURES: &[&str] = &[ + "fixtures/lang/eng_Latn.txt", + "fixtures/lang/cmn_Hani.txt", + "fixtures/modalities/code_mixed.txt", +]; + +/// The words `fixture_bench` injects as `normalized:true` added tokens, so the probe +/// exercises the same second-scan state the comparative benchmark runs under. +const NORMALIZED_WORDS: &[&str] = &["widgetron", "flibberjast", "zorptastic", "quibblenaut"]; + +fn median(mut samples: Vec) -> f64 { + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + samples[samples.len() / 2] +} + +fn timed(mut run: impl FnMut()) -> f64 { + run(); // warm-up + let mut samples = Vec::with_capacity(REPS); + for _ in 0..REPS { + let start = Instant::now(); + run(); + samples.push(start.elapsed().as_secs_f64()); + } + median(samples) +} + +fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String]) -> f64 { + let mut out = Vec::new(); + let mut pre_tokens = Vec::new(); + let mut scratch = pipeline.get_model().init_scratch(); + timed(|| { + for chunk in chunks { + out.clear(); + let _ = pipeline.encode_generic::( + chunk, + true, + &mut pre_tokens, + &mut scratch, + &mut out, + ); + black_box(&out); + black_box(&pre_tokens); + } + }) +} + +/// Whole lines accumulated up to `chunk_bytes`, `MAX_BYTES` in total. `chunk_bytes = 0` +/// keeps each line its own chunk (the short-input regime). +fn chunks_of(text: &str, chunk_bytes: usize) -> Vec { + let mut chunks = Vec::new(); + let mut current = String::new(); + let mut total = 0usize; + for line in text.lines().filter(|l| !l.trim().is_empty()) { + current.push_str(line); + if current.len() >= chunk_bytes { + total += current.len(); + chunks.push(std::mem::take(&mut current)); + if total >= MAX_BYTES { + return chunks; + } + } else { + current.push(' '); + } + } + if !current.is_empty() { + chunks.push(current); + } + chunks +} + +struct Ladder { + added: f64, + norm: f64, + split: f64, + model: f64, + post: f64, + total_mbs: f64, +} + +fn ladder(pipeline: &PipelineTokenizer, chunks: &[String], bytes: usize) -> Ladder { + let t_frame = stage_secs::<{ PipelineTokenizer::STAGE_FRAME }>(pipeline, chunks); + let t_norm = stage_secs::<{ PipelineTokenizer::STAGE_NORMALIZE }>(pipeline, chunks); + let t_split = stage_secs::<{ PipelineTokenizer::STAGE_SPLIT }>(pipeline, chunks); + let t_model = stage_secs::<{ PipelineTokenizer::STAGE_MODEL }>(pipeline, chunks); + let t_post = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(pipeline, chunks); + let nspb = |secs: f64| secs * 1e9 / bytes as f64; + Ladder { + added: nspb(t_frame.max(0.0)), + norm: nspb((t_norm - t_frame).max(0.0)), + split: nspb((t_split - t_norm).max(0.0)), + model: nspb((t_model - t_split).max(0.0)), + post: nspb((t_post - t_model).max(0.0)), + total_mbs: bytes as f64 / t_post / 1e6, + } +} + +/// The fused normalizer's pass, split into its three costs over the same chunks: +/// the space count, the `String` allocation, and the full rewrite (count + alloc + +/// write). All three follow `MetaspaceNormalizer::normalize`'s non-`drop_whitespace` +/// arm byte for byte, so `rewrite` should land on the ladder's norm marginal. +fn decompose(chunks: &[String], bytes: usize, prepend: bool) -> (f64, f64, f64) { + let space = Literal::new(b" ").unwrap(); + let delimiter = "\u{2581}"; + let counts: Vec = chunks + .iter() + .map(|c| space.count_matches(c.as_bytes())) + .collect(); + let nspb = |secs: f64| secs * 1e9 / bytes as f64; + + let count_only = timed(|| { + for chunk in chunks { + black_box(space.count_matches(chunk.as_bytes())); + } + }); + let alloc_only = timed(|| { + for (chunk, &count) in chunks.iter().zip(&counts) { + let s = String::with_capacity(chunk.len() + 2 * count + if prepend { 3 } else { 0 }); + black_box(&s); + } + }); + let rewrite = timed(|| { + for chunk in chunks { + let count = space.count_matches(chunk.as_bytes()); + if !prepend && count == 0 { + black_box(chunk.as_str()); + continue; + } + let mut rewritten = + String::with_capacity(chunk.len() + 2 * count + if prepend { 3 } else { 0 }); + if prepend { + rewritten.push_str(delimiter); + } + let mut prev = 0; + space.for_each_match(chunk.as_bytes(), |start| { + rewritten.push_str(&chunk[prev..start]); + rewritten.push_str(delimiter); + prev = start + 1; + }); + rewritten.push_str(&chunk[prev..]); + black_box(&rewritten); + } + }); + (nspb(count_only), nspb(alloc_only), nspb(rewrite)) +} + +fn main() { + let fixtures: Vec<(String, String)> = FIXTURES + .iter() + .map(|rel| { + let path = Path::new(DATA).join(rel); + let name = path.file_stem().unwrap().to_str().unwrap().to_string(); + (name, std::fs::read_to_string(&path).unwrap()) + }) + .collect(); + + for &(file, replica_prepend) in MODELS { + let path = Path::new(DATA).join(file); + let mut tok = match Tokenizer::from_file(&path) { + Ok(t) => t, + Err(e) => { + println!("== {file}: load failed: {e}"); + continue; + } + }; + let mut pipeline = match PipelineTokenizer::try_from(&tok) { + Ok(p) => p, + Err(e) => { + println!("== {file}: no pipeline: {e}"); + continue; + } + }; + // The ladder times the written stages; the zero-copy path gets its own A/B below. + pipeline.disable_zero_copy(); + let pretok = format!("{:?}", pipeline.get_pre_tokenizer()); + let pretok = pretok.split(['(', ' ']).next().unwrap_or("?"); + println!("== {file} (pre-tokenizer: {pretok})"); + + for (name, text) in &fixtures { + for (regime, chunk_bytes) in [("10kB", CHUNK_BYTES), ("line", 0)] { + let chunks = chunks_of(text, chunk_bytes); + let bytes: usize = chunks.iter().map(String::len).sum(); + let l = ladder(&pipeline, &chunks, bytes); + let total = l.added + l.norm + l.split + l.model + l.post; + println!( + " {name:<12} {regime:<4} ({:>5} chunks, {:>4} kB): added {:.3} | norm {:.3} | split {:.3} | model {:.3} | post {:.3} ns/B e2e {:.0} MB/s norm = {:.1}% of encode", + chunks.len(), + bytes / 1024, + l.added, + l.norm, + l.split, + l.model, + l.post, + l.total_mbs, + 100.0 * l.norm / total.max(1e-9), + ); + if let Some(prepend) = replica_prepend { + let (count, alloc, rewrite) = decompose(&chunks, bytes, prepend); + println!( + " {:12} {regime:<4} norm decomposed: count {count:.3} + alloc {alloc:.3} + write {:.3} = replica {rewrite:.3} (ladder said {:.3})", + "", + (rewrite - count - alloc).max(0.0), + l.norm, + ); + } + } + } + + // Zero-copy A/B: the same binary and corpus, the two paths interleaved rep by rep so + // frequency drift hits both equally. Ids are compared over the whole corpus first. + if replica_prepend.is_some() { + let zero_copy = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + zero_copy.has_zero_copy(), + "{file}: the zero-copy path should fire" + ); + let mut written = PipelineTokenizer::try_from(&tok).unwrap(); + written.disable_zero_copy(); + for (name, text) in &fixtures { + for (regime, chunk_bytes) in [("10kB", CHUNK_BYTES), ("line", 0)] { + let chunks = chunks_of(text, chunk_bytes); + let bytes: usize = chunks.iter().map(String::len).sum(); + for chunk in &chunks { + let ids = |p: &PipelineTokenizer| -> Vec { + p.encode(chunk, true) + .unwrap() + .iter() + .map(|t| t.id) + .collect() + }; + assert_eq!(ids(&zero_copy), ids(&written), "{file} {name}: ids diverge"); + } + let pass = |p: &PipelineTokenizer| { + let mut n = 0usize; + for chunk in &chunks { + n += p.encode(chunk, true).unwrap().len(); + } + black_box(n); + }; + pass(&zero_copy); // warm-up + pass(&written); + let mut zc = Vec::with_capacity(REPS); + let mut wr = Vec::with_capacity(REPS); + for _ in 0..REPS { + let t = Instant::now(); + pass(&zero_copy); + zc.push(t.elapsed().as_secs_f64()); + let t = Instant::now(); + pass(&written); + wr.push(t.elapsed().as_secs_f64()); + } + let (zc, wr) = (median(zc), median(wr)); + println!( + " A/B {name:<12} {regime:<4}: zero-copy {:.0} MB/s vs written {:.0} MB/s ({:+.1}%)", + bytes as f64 / zc / 1e6, + bytes as f64 / wr / 1e6, + 100.0 * (wr / zc - 1.0), + ); + } + } + } + + // t5 and albert declare more normalizers than the space rewrite (t5 a Precompiled + // charsmap, albert five steps and then one). The design only removes the rewrite, + // so its share is measured by stripping the declared normalizer: what is left in + // the norm marginal comes from the Metaspace pre-tokenizer's rewriting half. + if replica_prepend.is_none() { + let mut stripped = Tokenizer::from_file(&path).unwrap(); + let _ = stripped.with_normalizer(None::); + if let Ok(metaspace_only) = PipelineTokenizer::try_from(&stripped) { + let (name, text) = &fixtures[0]; + let chunks = chunks_of(text, CHUNK_BYTES); + let bytes: usize = chunks.iter().map(String::len).sum(); + let l = ladder(&metaspace_only, &chunks, bytes); + println!( + " declared normalizer stripped, {name} 10kB: norm {:.3} ns/B is the metaspace share", + l.norm, + ); + } + } + + // The second scan runs over every normalized chunk, but on an empty normalized + // vocabulary `Buckets::match_bytes` returns before touching the text. Injecting + // normalized added tokens (as `fixture_bench` does for every model) makes it a + // real pass; the frame marginal shows what the design's build-time gate saves. + let injected: Vec = NORMALIZED_WORDS + .iter() + .map(|w| AddedToken::from(*w, false).normalized(true)) + .collect(); + let _ = tok.add_tokens(injected); + if let Ok(with_normalized) = PipelineTokenizer::try_from(&tok) { + let (name, text) = &fixtures[0]; + let chunks = chunks_of(text, CHUNK_BYTES); + let bytes: usize = chunks.iter().map(String::len).sum(); + let before = ladder(&pipeline, &chunks, bytes); + let after = ladder(&with_normalized, &chunks, bytes); + println!( + " 2nd-scan probe on {name} 10kB: added {:.3} -> {:.3} ns/B with {} normalized added tokens", + before.added, + after.added, + NORMALIZED_WORDS.len(), + ); + } + println!(); + } +} diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 8fdec30d5e..a05fbc6dd0 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -776,6 +776,17 @@ impl PipelineBPE { self.ignore_merges } + /// Does this model seed the merge loop with characters (rather than raw bytes)? Only + /// character atoms can read a span through a [`SpaceSwap`], which swaps one char. + pub(crate) fn char_atoms(&self) -> bool { + matches!(self.atoms, Atoms::Chars { .. }) + } + + /// The id of the vocabulary piece spelled exactly `bytes`, if there is one. + pub(crate) fn id_of_bytes(&self, bytes: &[u8]) -> Option { + self.vocab.get_bytes(bytes) + } + /// Every piece of the vocabulary, as bytes. Used at build time to work out how the text may be /// cut into words; not cheap, so call it once. pub(crate) fn vocab_bytes(&self) -> Vec<(Vec, u32)> { @@ -785,6 +796,7 @@ impl PipelineBPE { fn merge_word( &self, sequence: &str, + swap: Option, merge_queue: &mut QuaternaryHeap, skip: &mut Vec, word: &mut Word, @@ -828,6 +840,14 @@ impl PipelineBPE { .char_indices() .map(|(i, c)| &sequence[i..i + c.len_utf8()]) { + // A raw `from` stands for its replacement: its atom is the replacement's + // id, as wide as what the rewrite would have written. + if let Some(swap) = swap + && char_str.as_bytes() == [swap.from] + { + word.add(swap.id, swap.len as usize); + continue; + } let char_len = char_str.len(); if let Some(char_id) = self.vocab.token_to_id(char_str) { word.add(char_id, char_len); @@ -856,6 +876,19 @@ impl PipelineBPE { } } +/// One raw character read as the one a pending rewrite would swap in: the raw `from` (a +/// single byte, all its route supports), the replacement's vocabulary id, and the byte width +/// the rewritten form would have. +/// +/// The zero-copy metaspace path (see [`crate::pre_tokenizers::proven_cuts`]) hands this to +/// [`PipelineBPE::tokenize_swapped`] so the merge loop seeds spans that were never rewritten. +#[derive(Debug, Clone, Copy)] +pub(crate) struct CharSwap { + pub(crate) from: u8, + pub(crate) id: u32, + pub(crate) len: u8, +} + impl PipelineBPE { /// One pre-token through the cache and the merge loop. This is the whole model /// step for a span and it cannot fail; it stands alone so the fused byte-level @@ -895,7 +928,7 @@ impl PipelineBPE { { output.push(PipelineToken { id }); } else { - self.merge_word(sequence, merge_queue, skip, word); + self.merge_word(sequence, None, merge_queue, skip, word); output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); } // The ids come back out of `output` because that is the only place both @@ -906,6 +939,49 @@ impl PipelineBPE { cache.insert(at, output[start..].iter().map(|token| token.id)); } } + + /// [`Self::tokenize_span`] for a raw span read through `swap`. The cache key is the + /// span's raw bytes; a rewritten form never holds the swap's `from`, so raw and rewritten + /// keys can only collide when they spell the same word, and the cached ids agree. + /// `ignore_merges` plays no part here: the zero-copy path is only built for models + /// without it, as the proven cuts it rides on already require. + pub(crate) fn tokenize_swapped( + &self, + sequence: &str, + swap: CharSwap, + scratch: &mut BpeScratch, + output: &mut Vec, + ) -> Result<()> { + if sequence.is_empty() { + return Ok(()); + } + let BpeScratch { + merge_queue, + skip, + word, + word_cache, + } = scratch; + + let mut placement = None; + if let Some(cache) = word_cache.as_mut() { + match cache.lookup(sequence.as_bytes()) { + Lookup::Hit(ids) => { + output.extend(ids.iter().map(|&id| PipelineToken { id })); + return Ok(()); + } + Lookup::Miss(at) => placement = at, + } + } + let start = output.len(); + self.merge_word(sequence, Some(swap), merge_queue, skip, word); + output.extend(word.get_chars_iter().map(|id| PipelineToken { id })); + if let Some(cache) = word_cache.as_mut() + && let Some(at) = placement + { + cache.insert(at, output[start..].iter().map(|token| token.id)); + } + Ok(()) + } } impl pipeline::Model for PipelineBPE { diff --git a/tokenizers/tk-encode/src/normalizers/metaspace.rs b/tokenizers/tk-encode/src/normalizers/metaspace.rs index 193c10d835..b54e79a21b 100644 --- a/tokenizers/tk-encode/src/normalizers/metaspace.rs +++ b/tokenizers/tk-encode/src/normalizers/metaspace.rs @@ -21,11 +21,10 @@ use std::borrow::Cow; -use atomsplit::literal::Literal; - use crate::normalizers::NormalizerWrapper; use crate::normalizers::replace::{Replace, ReplacePattern}; use crate::pre_tokenizers::whitespace::WhitespaceSplit; +use crate::tokenizer::pipeline::PendingRewrite; use crate::tokenizer::{Result, pipeline}; /// When [`MetaspaceNormalizer`] writes a delimiter at the start of the text it is given. @@ -66,6 +65,16 @@ impl MetaspaceNormalizer { } } + /// The swap and prepend as a [`PendingRewrite`]. The whitespace-dropping form deletes + /// characters, which no character-for-character rewrite can express, so it has none. + pub(crate) fn as_rewrite(&self) -> Option { + (!self.drop_whitespace).then_some(PendingRewrite { + from: ' ', + to: self.delimiter, + prepend: self.prepend, + }) + } + /// The one-pass stand-in for the leading `steps`, when they spell out this normalizer's job: /// `Prepend(c)` followed by `Replace(" " -> c)` (llama-2's normalizer), or a `Replace(" " -> c)` /// on its own (gemma-4's). Returns the stand-in and how many steps it covers; `None` when the @@ -109,73 +118,43 @@ fn space_swap(replace: &Replace) -> Option { } impl pipeline::Normalizer for MetaspaceNormalizer { + fn pending_rewrite(&self) -> Option { + self.as_rewrite() + } + fn normalize<'a>(&self, input: &'a str) -> Result> { + // Only spaces become delimiters; tabs and newlines are left alone. The swap and the + // prepend are exactly a [`PendingRewrite`], written here. + if let Some(rewrite) = self.as_rewrite() { + return rewrite.write(input); + } // Return empty input as is if input.is_empty() { return Ok(Cow::Borrowed(input)); } - if self.drop_whitespace { - // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the - // words back one after the other, each with its own delimiter. - let mut words = Vec::new(); - pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?; - // Exact when every word takes a delimiter; a word that already starts with one - // (`IfMissing`) leaves a few bytes spare. - let words_len: usize = words.iter().map(|span| span.range().len()).sum(); - let mut rewritten = - String::with_capacity(words_len + words.len() * self.delimiter.len_utf8()); - for span in &words { - let word = &input[span.range()]; - let prepend = match self.prepend { - PrependMode::Never => false, - // The text may already hold delimiters of its own: never write a second one. - PrependMode::IfMissing => !word.starts_with(self.delimiter), - PrependMode::Unconditional => true, - }; - if prepend { - rewritten.push(self.delimiter); - } - rewritten.push_str(word); - } - Ok(Cow::Owned(rewritten)) - } else { + // Whitespace is thrown away, so cut the text where `WhitespaceSplit` would and write the + // words back one after the other, each with its own delimiter. + let mut words = Vec::new(); + pipeline::PreTokenizer::pre_tokenize(&WhitespaceSplit, input, &mut words)?; + // Exact when every word takes a delimiter; a word that already starts with one + // (`IfMissing`) leaves a few bytes spare. + let words_len: usize = words.iter().map(|span| span.range().len()).sum(); + let mut rewritten = + String::with_capacity(words_len + words.len() * self.delimiter.len_utf8()); + for span in &words { + let word = &input[span.range()]; let prepend = match self.prepend { PrependMode::Never => false, - // A leading space counts as already marked: the swap turns it into a delimiter. - PrependMode::IfMissing => { - !input.starts_with(' ') && !input.starts_with(self.delimiter) - } + // The text may already hold delimiters of its own: never write a second one. + PrependMode::IfMissing => !word.starts_with(self.delimiter), PrependMode::Unconditional => true, }; - // Only spaces become delimiters; tabs and newlines are left alone. Counting them - // first sizes the rewrite exactly (a space is one byte, the delimiter up to four) - // and lets the swap stream through the batch scan instead of restarting a search - // at every space. - let space = Literal::new(b" ").expect("a space is not empty"); - let count = space.count_matches(input.as_bytes()); - // Nothing to prepend and nothing to swap: hand the input back instead of copying it. - if !prepend && count == 0 { - return Ok(Cow::Borrowed(input)); - } - let mut buf = [0u8; 4]; - let delimiter = self.delimiter.encode_utf8(&mut buf); - let mut rewritten = String::with_capacity( - input.len() - + (delimiter.len() - 1) * count - + if prepend { delimiter.len() } else { 0 }, - ); if prepend { - rewritten.push_str(delimiter); + rewritten.push(self.delimiter); } - let mut prev = 0; - space.for_each_match(input.as_bytes(), |start| { - rewritten.push_str(&input[prev..start]); - rewritten.push_str(delimiter); - prev = start + 1; - }); - rewritten.push_str(&input[prev..]); - Ok(Cow::Owned(rewritten)) + rewritten.push_str(word); } + Ok(Cow::Owned(rewritten)) } } diff --git a/tokenizers/tk-encode/src/normalizers/mod.rs b/tokenizers/tk-encode/src/normalizers/mod.rs index c16751ff47..effe8d888d 100644 --- a/tokenizers/tk-encode/src/normalizers/mod.rs +++ b/tokenizers/tk-encode/src/normalizers/mod.rs @@ -237,6 +237,25 @@ impl pipeline::Normalizer for NormalizerWrapper { Self::ByteLevel(bl) => pipeline::Normalizer::normalize(bl, input), } } + + fn pending_rewrite(&self) -> Option { + match self { + Self::BertNormalizer(bn) => pipeline::Normalizer::pending_rewrite(bn), + Self::StripNormalizer(sn) => pipeline::Normalizer::pending_rewrite(sn), + Self::StripAccents(sn) => pipeline::Normalizer::pending_rewrite(sn), + Self::NFC(nfc) => pipeline::Normalizer::pending_rewrite(nfc), + Self::NFD(nfd) => pipeline::Normalizer::pending_rewrite(nfd), + Self::NFKC(nfkc) => pipeline::Normalizer::pending_rewrite(nfkc), + Self::NFKD(nfkd) => pipeline::Normalizer::pending_rewrite(nfkd), + Self::Sequence(sequence) => pipeline::Normalizer::pending_rewrite(sequence), + Self::Lowercase(lc) => pipeline::Normalizer::pending_rewrite(lc), + Self::Nmt(nmt) => pipeline::Normalizer::pending_rewrite(nmt), + Self::Precompiled(pc) => pipeline::Normalizer::pending_rewrite(pc), + Self::Replace(rp) => pipeline::Normalizer::pending_rewrite(rp), + Self::Prepend(pp) => pipeline::Normalizer::pending_rewrite(pp), + Self::ByteLevel(bl) => pipeline::Normalizer::pending_rewrite(bl), + } + } } #[cfg(test)] diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index 5f2ea686ce..872a64308d 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; +use crate::normalizers::metaspace::PrependMode; use crate::pipeline; use crate::tokenizer::Decoder; use crate::tokenizer::pattern::Pattern; @@ -146,6 +147,23 @@ fn replace_matches<'a>( } impl pipeline::Normalizer for Replace { + /// A swap of one character for another qualifies. Regex patterns, multi-character + /// literals and deletions change the text's structure and have to be written. + fn pending_rewrite(&self) -> Option { + let ReplacePattern::String(pattern) = &self.pattern else { + return None; + }; + let mut pattern = pattern.chars(); + let from = pattern.next().filter(|_| pattern.next().is_none())?; + let mut content = self.content.chars(); + let to = content.next().filter(|_| content.next().is_none())?; + Some(pipeline::PendingRewrite { + from, + to, + prepend: PrependMode::Never, + }) + } + fn normalize<'a>(&self, input: &'a str) -> Result> { Ok(match &self.search { Search::Literal(literal) => { diff --git a/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs index 2ad2ee4047..0e5dcd443c 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/proven_cuts.rs @@ -26,12 +26,13 @@ use atomsplit::literal::Literal; +use crate::models::bpe::{CharSwap, PipelineBPE}; use crate::normalizers::NormalizerWrapper; use crate::normalizers::replace::{Replace, ReplacePattern}; use crate::pre_tokenizers::PreTokenizerWrapper; use crate::pre_tokenizers::split::SplitPattern; use crate::tokenizer::Result; -use crate::tokenizer::pipeline::{self, PipelineModel, Span}; +use crate::tokenizer::pipeline::{self, PendingRewrite, PipelineModel, Span}; /// Splits `▁`-spelled text into words, at the delimiters the vocabulary allows. See the module docs. #[derive(Debug, Clone)] @@ -162,6 +163,120 @@ fn space_replacement(replace: &Replace) -> Option { content.next().is_none().then_some(delimiter) } +/// [`ProvenCuts`] without the rewrite it runs on: the same cuts, found on the raw text. +/// +/// The tokenizers this serves normalize with a [`PendingRewrite`], swapping every space for +/// the delimiter (llama-2 prepends one too), so their rewritten text differs from the raw +/// text one character at a time. Every position [`ProvenCuts`] would cut is therefore +/// visible in the raw text: it is a raw `from` or a raw `to`. Cutting there directly means +/// the rewrite is never written; the model reads each raw span through a [`CharSwap`] +/// ([`PipelineBPE::tokenize_swapped`]), and only the one word a prepend touches is rewritten +/// for real. +/// +/// The veto is taken at its prefilter's word: a cut whose preceding byte could open a veto +/// piece is skipped without checking the piece itself. Skipping a cut never changes the ids +/// (each cut is only ever a speed-up), it only hands the model a longer span. The full check +/// reads the rewritten bytes around the cut, which is exactly what this path avoids building. +#[derive(Debug, Clone)] +pub(crate) struct ZeroCopyMetaspace { + /// The rewrite this path stands in for; also writes the one span the prepend touches. + rewrite: PendingRewrite, + /// The rewrite's `from`, as the single byte it encodes to. The cuts are only proven for + /// the space swap, so `from` is always one byte; comparing a byte (not a runtime-length + /// slice, which compiles to a `memcmp` call) is what keeps the scan tight. + from: u8, + /// UTF-8 bytes of the delimiter (the rewrite's `to`); only the first `to_len` are meaningful. + to: [u8; 4], + to_len: u8, + /// The delimiter's vocabulary id, seeded for every raw `from`. + to_id: u32, + /// [`Veto::bytes_before`]: last bytes of the pieces' `before` halves, as a 256-bit set. + veto_bytes_before: [u64; 4], +} + +impl ZeroCopyMetaspace { + /// `None` when the raw text cannot stand in for the rewritten text: the cuts were proven + /// for the space swap only, and byte-atom models spell text in another alphabet. + pub(crate) fn build( + rewrite: PendingRewrite, + cuts: &ProvenCuts, + bpe: &PipelineBPE, + ) -> Option { + if rewrite.from != ' ' || !bpe.char_atoms() { + return None; + } + let mut to = [0u8; 4]; + let encoded = rewrite.to.encode_utf8(&mut to); + if cuts.delimiter.pattern() != encoded.as_bytes() { + return None; + } + let to_id = bpe.id_of_bytes(encoded.as_bytes())?; + let to_len = encoded.len() as u8; + Some(Self { + rewrite, + from: rewrite.from as u8, + to, + to_len, + to_id, + veto_bytes_before: cuts.veto.bytes_before, + }) + } + + pub(crate) fn rewrite(&self) -> PendingRewrite { + self.rewrite + } + + pub(crate) fn swap(&self) -> CharSwap { + CharSwap { + from: self.from, + id: self.to_id, + len: self.to_len, + } + } + + /// Cuts `text` where [`ProvenCuts`] would cut its rewritten form, as raw byte offsets. + /// Returns whether the first span takes the prepended delimiter; the caller writes that + /// one span's rewrite for real and reads the rest through the swap. + pub(crate) fn cut(&self, text: &str, out: &mut Vec) -> bool { + let bytes = text.as_bytes(); + let to = &self.to[..self.to_len as usize]; + let mut start = 0usize; + let mut search = 0usize; + while let Some(found) = memchr::memchr2(self.from, to[0], &bytes[search..]) { + let at = search + found; + let width = if bytes[at] == self.from { + 1 + } else if bytes[at..].starts_with(to) { + to.len() + } else { + // The delimiter's first byte opens other characters too; not a delimiter. + search = at + 1; + continue; + }; + // The same three refusals as [`ProvenCuts::pre_tokenize`]: the word would be + // empty, the delimiter sits inside a group (its predecessor is a raw `from` or + // a raw `to`), or the veto's prefilter byte fires. The predecessor byte is + // never rewritten (a rewritten one is the group case), so raw is exact here. + if at != start + && !(bytes[at - 1] == self.from || bytes[..at].ends_with(to)) + && !self.prefilter_forbids(bytes[at - 1]) + { + out.push(Span::new(start as u32, at as u32)); + start = at; + } + search = at + width; + } + if start < bytes.len() { + out.push(Span::new(start as u32, bytes.len() as u32)); + } + self.rewrite.prepends(text) + } + + fn prefilter_forbids(&self, previous: u8) -> bool { + self.veto_bytes_before[(previous >> 6) as usize] & (1 << (previous & 63)) != 0 + } +} + /// How many veto pieces we put up with before giving up on cutting at all. const MAX_VETO_PIECES: usize = 32; diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 92bd56de48..a43fc2048c 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -19,14 +19,17 @@ use crate::vocab::bucket_added_vocabulary::{ }; use crate::{ ModelWrapper, PostProcessorWrapper, PreTokenizerWrapper, Token, Tokenizer, - normalizers::{NormalizerWrapper, metaspace::MetaspaceNormalizer}, + normalizers::{ + NormalizerWrapper, + metaspace::{MetaspaceNormalizer, PrependMode}, + }, pre_tokenizers::{ bert::BertPreTokenizer, delimiter::CharDelimiterSplit, digits::Digits, fixed_length::FixedLength, metaspace, - proven_cuts::{self, ProvenCuts}, + proven_cuts::{self, ProvenCuts, ZeroCopyMetaspace}, punctuation::Punctuation, sequence::PipelineSequence, split::{Split as SplitPretok, SplitPattern}, @@ -66,15 +69,160 @@ pub(crate) fn classify_into_spans( pub trait Normalizer { fn normalize<'a>(&self, input: &'a str) -> Result>; + + /// The rewrite this normalizer could leave pending instead of writing, or `None` when its + /// rewrite restructures the text and has to be written. Only rewrites that change the + /// text one character at a time qualify — see [`PendingRewrite`]. Today those are + /// [`MetaspaceNormalizer`] (the SentencePiece space swap) and a single-character literal + /// [`Replace`](crate::normalizers::replace::Replace). + fn pending_rewrite(&self) -> Option { + None + } } -/// Runs `normalizers` in order, each one seeing what the one before it produced. +/// A rewrite that changes the text one character at a time: every `from` reads as `to`, and +/// `prepend` may put one `to` in front of the chunk. /// -/// A normalizer returns a [`Cow`] (copy-on-write): a borrow when it had nothing to change, an owned -/// `String` when it rewrote the text. +/// Because positions map one to one, such a rewrite does not have to be written into a copy: +/// a consumer can read the rewritten form off the raw text. [`NormalizedText`] carries one +/// left pending by the normalizer chain, and the zero-copy encode path consumes it in place; +/// [`PendingRewrite::write`] produces the rewritten text for every other consumer. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct PendingRewrite { + pub(crate) from: char, + pub(crate) to: char, + pub(crate) prepend: PrependMode, +} + +impl PendingRewrite { + /// Whether the rewrite puts a `to` in front of `text`. Empty text is handed back as is, + /// so it never takes one. + pub(crate) fn prepends(&self, text: &str) -> bool { + !text.is_empty() + && match self.prepend { + PrependMode::Never => false, + // Text already opening with the rewrite's own characters counts as marked + // (a leading `from` becomes a `to` under the swap). + PrependMode::IfMissing => { + !text.starts_with(self.from) && !text.starts_with(self.to) + } + PrependMode::Unconditional => true, + } + } + + /// Writes the rewrite into a copy: one pass, one right-sized allocation, and a borrow + /// when there is nothing to swap or prepend. + pub(crate) fn write<'a>(&self, input: &'a str) -> Result> { + if input.is_empty() { + return Ok(Cow::Borrowed(input)); + } + let prepend = self.prepends(input); + let mut from_buf = [0u8; 4]; + let from = self.from.encode_utf8(&mut from_buf); + let mut to_buf = [0u8; 4]; + let to = self.to.encode_utf8(&mut to_buf); + // Counting the swaps first sizes the copy exactly and lets them stream through the + // batch scan instead of restarting a search at every match. + let pattern = + atomsplit::literal::Literal::new(from.as_bytes()).expect("a char is never empty"); + let count = pattern.count_matches(input.as_bytes()); + if !prepend && count == 0 { + return Ok(Cow::Borrowed(input)); + } + let mut rewritten = String::with_capacity( + input.len() - count * from.len() + + count * to.len() + + if prepend { to.len() } else { 0 }, + ); + if prepend { + rewritten.push_str(to); + } + let mut prev = 0; + pattern.for_each_match(input.as_bytes(), |start| { + rewritten.push_str(&input[prev..start]); + rewritten.push_str(to); + prev = start + from.len(); + }); + rewritten.push_str(&input[prev..]); + Ok(Cow::Owned(rewritten)) + } +} + +impl Normalizer for PendingRewrite { + fn normalize<'a>(&self, input: &'a str) -> Result> { + self.write(input) + } + + fn pending_rewrite(&self) -> Option { + Some(*self) + } +} + +/// What a normalizer chain produced for one chunk: the normalized text, or that text one +/// pending rewrite short of it. +/// +/// [`NormalizedText::from_chain`] leaves a trailing [`PendingRewrite`] unwritten; the +/// zero-copy encode path reads the rewritten form off `text`'s raw characters, and every +/// other consumer calls [`NormalizedText::write`] to get plain text. +pub struct NormalizedText<'a> { + text: Cow<'a, str>, + pending: Option, +} + +impl<'a> NormalizedText<'a> { + /// Runs `normalizers` over `input`, leaving a trailing pending rewrite unwritten. A + /// pending rewrite anywhere else in the chain runs as declared: the step after it needs + /// the written text. + pub(crate) fn from_chain(normalizers: &[N], input: &'a str) -> Result { + if let [head @ .., last] = normalizers + && let Some(pending) = last.pending_rewrite() + { + return Ok(Self { + text: normalize_all(head, input)?, + pending: Some(pending), + }); + } + Ok(Self { + text: normalize_all(normalizers, input)?, + pending: None, + }) + } + + /// The normalized text, with any pending rewrite written into it. + pub fn write(self) -> Result> { + match self.pending { + None => Ok(self.text), + Some(rewrite) => normalize_step(&rewrite, self.text), + } + } +} + +/// One normalizer step over the running text: the copy-on-write bookkeeping one link of +/// [`normalize_all`]'s chain needs. /// -/// Chaining them needs care, because an owned `String` produced halfway through the chain is local to this function. -/// The next normalizer may hand back a borrow of that locally owned [`String`], and that borrow cannot outlive the `String`. +/// A normalizer returns a [`Cow`] (copy-on-write): a borrow when it had nothing to change, an +/// owned `String` when it rewrote the text. Chaining needs care, because an owned `String` +/// produced halfway through the chain is local: the next normalizer may hand back a borrow of +/// that `String`, and that borrow cannot outlive it. +fn normalize_step<'a, N: Normalizer>(normalizer: &N, cow: Cow<'a, str>) -> Result> { + Ok(match cow { + // Still the caller's input, which outlives us: pass it straight on. + Cow::Borrowed(s) => normalizer.normalize(s)?, + Cow::Owned(s) => { + let out = match normalizer.normalize(&s)? { + // Rewritten again: keep the new `String`, drop ours. + Cow::Owned(o) => Some(o), + // Handed `s` back untouched: keep the `String` we already own. + Cow::Borrowed(b) if b.as_ptr() == s.as_ptr() && b.len() == s.len() => None, + // A borrow into `s` (for example, a substring of `s`): copy it out before `s` is dropped. + Cow::Borrowed(b) => Some(b.to_owned()), + }; + Cow::Owned(out.unwrap_or(s)) + } + }) +} + +/// Runs `normalizers` in order, each one seeing what the one before it produced. /// /// Text no normalizer touches is never copied: it stays a borrow of `input` throughout. pub(crate) fn normalize_all<'a, N: Normalizer>( @@ -83,21 +231,7 @@ pub(crate) fn normalize_all<'a, N: Normalizer>( ) -> Result> { let mut cow: Cow<'a, str> = Cow::Borrowed(input); for normalizer in normalizers { - cow = match cow { - // Still `input` itself, which outlives us: pass it straight on. - Cow::Borrowed(s) => normalizer.normalize(s)?, - Cow::Owned(s) => { - let out = match normalizer.normalize(&s)? { - // Rewritten again: keep the new `String`, drop ours. - Cow::Owned(o) => Some(o), - // Handed `s` back untouched: keep the `String` we already own. - Cow::Borrowed(b) if b.as_ptr() == s.as_ptr() && b.len() == s.len() => None, - // A borrow into `s` (for example, a substring of `s`): copy it out before `s` is dropped. - Cow::Borrowed(b) => Some(b.to_owned()), - }; - Cow::Owned(out.unwrap_or(s)) - } - }; + cow = normalize_step(normalizer, cow)?; } Ok(cow) } @@ -124,6 +258,13 @@ impl Normalizer for PipelineNormalizer { Self::Metaspace(normalizer) => normalizer.normalize(input), } } + + fn pending_rewrite(&self) -> Option { + match self { + Self::Declared(normalizer) => normalizer.pending_rewrite(), + Self::Metaspace(normalizer) => normalizer.pending_rewrite(), + } + } } /// Range-based pre-tokenization: yields spans into the input rather than owned @@ -479,6 +620,10 @@ pub struct PipelineTokenizer { model: PipelineModel, post_processor: PipelinePostProcessor, scratch_pool: ScratchPool, + /// When set, a full encode runs each chunk through [`ZeroCopyMetaspace`] instead of the + /// normalize/pre-tokenize stages: the rewrite those stages would write is read off the + /// raw text. `None` for every other pipeline shape; see [`PipelineTokenizer::try_from`]. + zero_copy: Option, } /// A pool of [`PipelineModelScratch`]. @@ -693,11 +838,26 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer { None => pre_tokenizer, }; + // A chain ending in a pending rewrite, followed by proven cuts: the rewrite can be + // read off the raw text instead of written (see [`ZeroCopyMetaspace`]). Added tokens + // flagged `normalized` are matched against rewritten text, which this path never + // builds, so one such token keeps the written stages. + let zero_copy = match (normalizers.last(), &pre_tokenizer, &model) { + (Some(last), PipelinePreTokenizer::ProvenCuts(cuts), PipelineModel::BPE(bpe)) + if !added_vocabulary.has_normalized_tokens() => + { + last.pending_rewrite() + .and_then(|rewrite| ZeroCopyMetaspace::build(rewrite, cuts, bpe)) + } + _ => None, + }; + Ok(Self { added_vocabulary, normalizers, pre_tokenizer, model, + zero_copy, post_processor: tok .get_post_processor() .map(PipelinePostProcessor::try_from) @@ -823,11 +983,31 @@ impl PipelineTokenizer { output.push(PipelineToken { id: token }); } Segment::Text(chunk) => { - let normalized: Cow = if STAGE >= Self::STAGE_NORMALIZE { - normalize_all(&self.normalizers, chunk)? + let normalized = if STAGE >= Self::STAGE_NORMALIZE { + NormalizedText::from_chain(&self.normalizers, chunk)? } else { - Cow::Borrowed(chunk) + NormalizedText { + text: Cow::Borrowed(chunk), + pending: None, + } }; + // Full encodes consume a pending rewrite in place when the pipeline has + // the zero-copy route for it; the partial stages write it below, so the + // ablation ladder times the stages the zero-copy path replaces. + if STAGE >= Self::STAGE_POSTPROCESS + && normalized.pending.is_some() + && let Some(zero_copy) = &self.zero_copy + { + self.encode_chunk_zero_copy( + zero_copy, + &normalized.text, + pre_tokens, + scratch, + output, + )?; + continue; + } + let normalized: Cow = normalized.write()?; // Extract special tokens from the normalized input for segment in @@ -933,6 +1113,49 @@ impl PipelineTokenizer { }); true } + + /// One chunk through [`ZeroCopyMetaspace`]: cut the raw text, then read each span through + /// the swap. The rewritten chunk is never built, and the special-token pass over it is + /// skipped with it — the path is only built when no added token matches normalized text. + fn encode_chunk_zero_copy( + &self, + zero_copy: &ZeroCopyMetaspace, + chunk: &str, + pre_tokens: &mut Vec, + scratch: &mut PipelineModelScratch, + output: &mut Vec, + ) -> Result<()> { + let (PipelineModel::BPE(bpe), PipelineModelScratch::BPE(scratch)) = (&self.model, scratch) + else { + unreachable!("`try_from` only builds `zero_copy` for a BPE model"); + }; + pre_tokens.clear(); + let prepend = zero_copy.cut(chunk, pre_tokens); + let mut spans = pre_tokens.iter(); + if prepend && let Some(first) = spans.next() { + // The prepend rewrites this one word for real, so its cache key is the rewritten + // form; every other span's key is its raw bytes. + let rewritten = zero_copy.rewrite().write(&chunk[first.range()])?; + bpe.tokenize_pipeline(&rewritten, scratch, output)?; + } + let swap = zero_copy.swap(); + for span in spans { + bpe.tokenize_swapped(&chunk[span.range()], swap, scratch, output)?; + } + Ok(()) + } + + /// Turns the zero-copy metaspace path off, so a full encode runs the written rewrite. + #[doc(hidden)] // public only so `examples/normalize_claims.rs` can A/B the two paths + pub fn disable_zero_copy(&mut self) { + self.zero_copy = None; + } + + /// Whether full encodes take the zero-copy metaspace path. + #[doc(hidden)] // public only so `examples/normalize_claims.rs` can verify the path fires + pub fn has_zero_copy(&self) -> bool { + self.zero_copy.is_some() + } } /// Streaming decoder over a [`PipelineTokenizer`]; see [`PipelineTokenizer::decode_stream`]. @@ -1491,6 +1714,250 @@ mod tests { assert_pipeline_matches_reference(&tok, "x a"); } + /// The metaspace-shaped BPE from the proven-cuts tests, plus an unk token and a `▁▁` + /// piece: the swap, the delimiter-group rule and the veto prefilter all have something + /// to act on. + fn metaspace_bpe_tokenizer( + normalizer_json: &str, + pre_tokenizer_json: Option<&str>, + ) -> Tokenizer { + use crate::models::bpe::{BpeBuilder, Merges, Vocab}; + + let vocab: Vocab = [ + ("", 0u32), + ("▁", 1), + ("<", 2), + (">", 3), + ("/", 4), + ("s", 5), + ("p", 6), + ("a", 7), + ("b", 8), + ("▁", 10), + (">▁", "▁"), + (">▁", "".to_string()) + .build() + .unwrap(); + let mut tok = Tokenizer::new(bpe); + tok.with_normalizer(Some( + serde_json::from_str::(normalizer_json).unwrap(), + )) + .unwrap(); + if let Some(json) = pre_tokenizer_json { + tok.with_pre_tokenizer(Some( + serde_json::from_str::(json).unwrap(), + )); + } + tok + } + + /// The two spellings of the space rewrite that reach the proven cuts, copied from the + /// files in `data/`. + const LLAMA_STYLE_NORMALIZER: &str = r#"{"type":"Sequence","normalizers":[{"type":"Prepend","prepend":"▁"},{"type":"Replace","pattern":{"String":" "},"content":"▁"}]}"#; + const GEMMA_STYLE_NORMALIZER: &str = + r#"{"type":"Replace","pattern":{"String":" "},"content":"▁"}"#; + const GEMMA_STYLE_PRE_TOKENIZER: &str = r#"{"type":"Split","pattern":{"String":" "},"behavior":"MergedWithPrevious","invert":false}"#; + + /// Chunk shapes the raw-coordinate cuts must reproduce: leading, trailing and grouped + /// spaces, raw delimiters, the veto text (cut by the written path, kept whole by the + /// prefilter), characters outside the vocabulary, and no spaces at all. + const ZERO_COPY_TEXTS: &[&str] = &[ + "sp a", + " sp", + "a b ", + " b a ", + "", + " ", + " ", + "sp", + "▁", + "▁a b", + "a▁b sp", + " ", + "a b", + "p a", + "a\tb\nsp", + "x y", + ]; + + /// Both space-rewrite spellings reach the fused-normalizer + proven-cuts shape, so both + /// must take the zero-copy path, and its ids must match the written rewrite's on every + /// chunk shape — including the veto text, where the two paths cut differently and only + /// the ids agree. + #[test] + fn zero_copy_ids_match_the_written_rewrite() { + for (name, normalizer, pre_tokenizer) in [ + ("llama-2 shape", LLAMA_STYLE_NORMALIZER, None), + ( + "gemma-4 shape", + GEMMA_STYLE_NORMALIZER, + Some(GEMMA_STYLE_PRE_TOKENIZER), + ), + ] { + let tok = metaspace_bpe_tokenizer(normalizer, pre_tokenizer); + let zero_copy = PipelineTokenizer::try_from(&tok).unwrap(); + assert!( + zero_copy.zero_copy.is_some(), + "{name} should take the zero-copy path" + ); + let mut written = PipelineTokenizer::try_from(&tok).unwrap(); + written.disable_zero_copy(); + for text in ZERO_COPY_TEXTS { + let got: Vec = zero_copy + .encode(text, false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + let want: Vec = written + .encode(text, false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(got, want, "{name}: {text:?}"); + assert_pipeline_matches_reference(&tok, text); + } + } + } + + /// An added token flagged `normalized` is matched against rewritten text, which the + /// zero-copy path never builds; one such token must keep the written stages. + #[test] + fn a_normalized_added_token_keeps_the_written_rewrite() { + let mut tok = metaspace_bpe_tokenizer(LLAMA_STYLE_NORMALIZER, None); + let _ = tok.add_tokens([crate::AddedToken::from("spa", false).normalized(true)]); + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + assert!(pipeline.zero_copy.is_none()); + assert_pipeline_matches_reference(&tok, "b spa sp a"); + } + + /// The chain hands a trailing pending rewrite back unwritten; `write` must produce + /// exactly what running the chain as declared produces. + #[test] + fn a_chain_ending_in_a_pending_rewrite_leaves_it_unwritten() { + let chain = [ + PipelineNormalizer::Declared(serde_json::from_str(r#"{"type":"Lowercase"}"#).unwrap()), + PipelineNormalizer::Metaspace(MetaspaceNormalizer::new( + '▁', + PrependMode::Unconditional, + false, + )), + ]; + for text in ["Hello World", "", " A ", "no_spaces", "▁Marked"] { + let normalized = NormalizedText::from_chain(&chain, text).unwrap(); + assert!(normalized.pending.is_some(), "{text:?}"); + assert_eq!( + NormalizedText::from_chain(&chain, text) + .unwrap() + .write() + .unwrap(), + normalize_all(&chain, text).unwrap(), + "{text:?}" + ); + } + } + + /// A pending rewrite anywhere but last runs as declared: the step after it needs the + /// written text. + #[test] + fn a_mid_chain_rewrite_is_written() { + let chain = [ + PipelineNormalizer::Metaspace(MetaspaceNormalizer::new( + '▁', + PrependMode::Unconditional, + false, + )), + PipelineNormalizer::Declared(serde_json::from_str(r#"{"type":"Lowercase"}"#).unwrap()), + ]; + let normalized = NormalizedText::from_chain(&chain, "Hello World").unwrap(); + assert!(normalized.pending.is_none()); + assert_eq!( + NormalizedText::from_chain(&chain, "Hello World") + .unwrap() + .write() + .unwrap(), + normalize_all(&chain, "Hello World").unwrap() + ); + } + + /// A `Replace` swapping one character for another is a pending rewrite too; anything + /// changing the text's structure is not. + #[test] + fn a_single_char_replace_is_a_pending_rewrite() { + let single: NormalizerWrapper = + serde_json::from_str(r#"{"type":"Replace","pattern":{"String":"x"},"content":"y"}"#) + .unwrap(); + assert_eq!( + Normalizer::pending_rewrite(&single), + Some(PendingRewrite { + from: 'x', + to: 'y', + prepend: PrependMode::Never, + }) + ); + let refused = [ + ( + "a multi-char pattern", + r#"{"type":"Replace","pattern":{"String":"xy"},"content":"y"}"#, + ), + ( + "multi-char content", + r#"{"type":"Replace","pattern":{"String":"x"},"content":"yy"}"#, + ), + ( + "a deletion", + r#"{"type":"Replace","pattern":{"String":"x"},"content":""}"#, + ), + ]; + for (name, json) in refused { + let replace: NormalizerWrapper = serde_json::from_str(json).unwrap(); + assert_eq!(Normalizer::pending_rewrite(&replace), None, "{name}"); + } + } + + /// Skipped when the fixtures have not been fetched. + #[test] + fn zero_copy_is_built_for_the_real_sentencepiece_configs() { + for file in ["llama-2.json", "gemma-4.json"] { + let path = format!("../data/{file}"); + if !std::path::Path::new(&path).exists() { + eprintln!("skip {file}: not present (fetch with `make bench-models`)"); + continue; + } + let tok = crate::Tokenizer::from_file(&path).unwrap(); + let pipeline = PipelineTokenizer::try_from(&tok).unwrap(); + assert!(pipeline.zero_copy.is_some(), "{file}"); + } + } + #[test] fn segment_iterator_yields_text_and_specials_in_order() { let input = "aabbcc"; @@ -1914,3 +2381,4 @@ mod tests { assert_eq!(cache.lookup(b"hello").hit(), Some(&[7u32][..])); } } + diff --git a/tokenizers/tk-encode/src/vocab/bucket_added_vocabulary.rs b/tokenizers/tk-encode/src/vocab/bucket_added_vocabulary.rs index 5bc0be6ddc..5dc13e080a 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_added_vocabulary.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_added_vocabulary.rs @@ -204,6 +204,12 @@ impl AddedVocabulary { self.vocab.is_empty() && self.normalized_vocab.is_empty() } + /// Does any added token match against normalized text? When none does, the pass over each + /// normalized chunk has nothing to find, and an encode path may skip building that text. + pub(crate) fn has_normalized_tokens(&self) -> bool { + !self.normalized_vocab.is_empty() + } + /// Get the additional vocabulary (union of both matchers; normalized tokens appear by their /// normalized form, since the original content isn't retained after partitioning). pub fn get_vocab(&self) -> AHashMap { From 972adc0f4364bbb31191e052dafad18a59adfe44 Mon Sep 17 00:00:00 2001 From: SBrandeis <33657802+SBrandeis@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:34:58 +0200 Subject: [PATCH 13/13] fix(bench): make the stage ladder measure the route encode actually takes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stage breakdown drove `encode_generic::`, which gated the fused scan at `STAGE_MODEL` and the zero-copy route at `STAGE_POSTPROCESS`. Below those the ladder fell back to the staged array FSM, so on every native-FSM model the `pre_tokenize` bar timed a splitter no shipping encode uses, and the `model` bar was a subtraction across two different pipelines (fused split+model minus unfused split), clamped at zero so the pathology never showed. The same held for the metaspace models: their whole zero-copy win landed in `post` or vanished into the clamp. `pretok_vs_regex` inherited the bad number and undersold the split against onig/fancy/pcre2/logos by roughly 2x. Both routes are now STAGE-aware, so every rung runs the route a full encode runs. A fused pipeline's split rung drives the same masked FSM scan the model rung does and emits each span into `pre_tokens`; a zero-copy pipeline's split rung is its proven-cut pass, with neither the rewrite write nor the added-token scan over rewritten text that the route never performs. That last one also fixes the frame rung, which was charging zero-copy models for a normalized added-token scan they never run. Measured on gpt2 (10 kB warm chunks, 5 reps, M3 Max), ns/byte: pre_tokenize eng 1.73 -> 0.57 swe 1.28 -> 0.61 cmn 1.23 -> 1.20 vs onig eng ~20x -> 63.8x Production throughput is unchanged: threading `pre_tokens` into the fused emit adds a closure capture that folds away in the `STAGE_MODEL` specialization. Alternating two binaries at process level over 5 rounds x 7 reps, 4 models x 2 fixtures, the spread is -1.6% to +2.5% with no systematic direction, inside this bench's noise floor. Also reported so the chart cannot be read as more complete than it is: - `route` per model, since the three decompose differently. - `clamped`, the residual the five stages do not account for. A rung at or below the noise floor goes negative and clamps, and a `clamped` far from zero means that fixture's breakdown should not be trusted. - `alloc`, the per-chunk output/span buffer allocation `encode` pays and the reused-buffer rungs do not. It measures 0.07 to 0.11 ns/byte on gpt2, so `Vec` growth is not the gap between the ladder's total and the throughput phase. Timing the real `encode` to close that gap was tried and dropped: phase 2 shares one pipeline across fixtures, its scratch pool saturates, and the delta came out at +0.9 to +7.2 ns/byte measuring cache eviction rather than call overhead. The remainder is left stated and unattributed instead. The split rung pays one span write per pre-token that a fused encode does not, worth 0.03 to 0.07 ns/byte (scan with a span push 0.638 vs with an xor 0.610 on gpt2/eng), and the model rung carries the checked `&str` slice, worth 0.15 to 0.27. Both are documented on `stage_secs`. Rendering the chart then turned up three readability defects the JSON hid. A row's stages sum past 100% wherever a rung clamped, and the bars were drawn against a fixed `PLOT_W`, so those rows ran off the right edge: ell_Grek at 108% and added_normalized_sparse at 109% were cut off, which reads as a broken chart rather than as noise. The plot is now scaled by the widest row, so an overflowing bar crosses the 100% gridline instead of leaving the canvas, and the subtitle carries the mean parts-sum. The header claimed "each bar = 100%", never quite true for the same reason, and now says 100% is that fixture's total. The new route note pushed the subtitle past the SVG width and clipped its tail, so stages under 2% mean go unlisted (on gpt2 that was "post 1% · normalize 0%"). Three tests pin the ladder, since `pipeline_oracle` only ever drives `STAGE_POSTPROCESS` and so proves ids without proving which path made them. The fused one asserts `encode_fused` returns true at `STAGE_SPLIT` rather than only comparing spans: the staged FSM produces identical spans, so span equality alone would not have caught this bug. Verified red by making the fused scan decline below `STAGE_MODEL`. 370 lib tests, `pipeline_oracle` 9/9 against released 0.23.1, clippy and fmt clean. `zero_copy_is_built_for_the_real_sentencepiece_configs` was already red at ac507686 and is untouched here. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HorT6V8fzJXUmdJwJPXfS5 --- .github/scripts/render_pipeline_bench.py | 63 +++- .../tk-encode/examples/fixture_bench.rs | 106 ++++++- .../tk-encode/src/tokenizer/pipeline.rs | 291 ++++++++++++++++-- 3 files changed, 408 insertions(+), 52 deletions(-) diff --git a/.github/scripts/render_pipeline_bench.py b/.github/scripts/render_pipeline_bench.py index ed7e92c716..41670db0b4 100644 --- a/.github/scripts/render_pipeline_bench.py +++ b/.github/scripts/render_pipeline_bench.py @@ -47,6 +47,12 @@ # `added_split` = added/special-token scan (AddedVocabulary), `pre_tokenize` = # pre-tokenizer split — two distinct splitting costs. `post` = special-token # id-frame splice (~0 for models with no post-processor). The five sum to `total`. +# +# `total` is the ladder's own full encode, NOT the public `encode` call the throughput +# phase times: that additionally acquires a pooled scratch and returns an owned Vec. +# The subtitle names that exclusion rather than the stack absorbing it, because the +# difference is part real cost and part pooled-vs-fresh cache state, and no measurement in +# the bench separates the two. STAGES = [("added_split", "added-token"), ("normalize", "normalize"), ("pre_tokenize", "pre-tokenize"), ("model", "model"), ("post", "post")] STAGE_INK = {"added_split": "#7a5ea8", "normalize": "#2a9d8f", @@ -640,12 +646,20 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label): sink = STAGE_INK rows = [r for r in model["results"] if "stage_ns_per_byte" in r] + # A row's stages are shares of that row's `total`, and they can sum past 100% when a + # rung measured at or below the noise floor and clamped to zero. Scaling the plot by the + # widest row keeps such a bar on the canvas and visibly past the 100% gridline, instead + # of running off the right edge as though the chart were broken. + span = max([1.0] + [sum(r["stage_ns_per_byte"].get(k, 0.0) for k, _ in STAGES) + / (r["stage_ns_per_byte"]["total"] or 1.0) for r in rows]) + scale = PLOT_W / span + top = 84 col_x = GUTTER + PLOT_W + PAD_R + COL_W - 16 body = [f'total ns/B · ×speedup', f'' - f'share of pipeline encode time (each bar = 100%) · label = share% · ns/B'] + f'share of pipeline encode time (100% = that fixture\'s total) · label = share% · ns/B'] y = top for key, title in GROUPS: group_rows = sorted((r for r in rows if r["group"] == key), @@ -665,7 +679,7 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label): for skey, _ in STAGES: val = s.get(skey, 0.0) frac = val / total - seg = frac * PLOT_W # each bar fills PLOT_W (100% of this fixture's total) + seg = frac * scale # 100% of this fixture's total is `scale` wide if seg > 0.4: body.append(f'') @@ -689,7 +703,7 @@ def stage_chart_svg(model, subtitle_base, meta, baseline_label): grid = [] for pct in (0, 25, 50, 75, 100): - gx = GUTTER + pct / 100 * PLOT_W + gx = GUTTER + pct / 100 * scale grid.append(f'') grid.append(f'= 0.02) + # The route is part of the reading, not decoration: each one decomposes into stages + # differently, and the split bar only means the shipped splitter because every rung + # drives this route. Older JSON has no `route` key, so say so rather than guess. + route = model.get("route") or "route not recorded" + subtitle = f'{model["shape"]} · {route} route · {mix_txt}' + # Two things the reader would otherwise have to infer. The rungs are independent + # medians, so one at or below the noise floor clamps to zero and the parts then sum + # past 100%, which is visible as a bar overrunning the axis. And `total` is the + # ladder's own encode, without the pooled scratch and returned Vec of a real call. + sums = [sum(r["stage_ns_per_byte"].get(k, 0.0) for k, _ in STAGES) + / r["stage_ns_per_byte"]["total"] + for r in rows if r["stage_ns_per_byte"].get("total")] + if sums: + subtitle += f' · parts sum {100 * sum(sums) / len(sums):.0f}% (rung noise)' + if any("alloc" in r["stage_ns_per_byte"] for r in rows): + subtitle += " · excludes encode() call overhead" return svg_doc(ink, height, f'{model["model"]} — Pipeline encode stage mix', subtitle, "".join(grid) + "".join(body) + legend, meta, subtitle_base) @@ -884,9 +915,9 @@ def x(v): def pretok_compare_md(model): - """Per-fixture 'classify + fsm vs a regex engine' table — the pre-tokenize split - vs onig/fancy/pcre2/logos on the model's own regex, WITH and WITHOUT SIMD. - Rendered only for regex pre-tokenizers (a reference is non-null).""" + """Per-fixture table of the shipped pre-tokenize split against onig, fancy, pcre2 and + logos on the model's own regex. Rendered only for regex pre-tokenizers (a reference is + non-null).""" engines = ("onig", "fancy", "pcre2", "logos") def has_ref(r): @@ -898,13 +929,15 @@ def has_ref(r): cell = lambda v: fnum(v, "{:.1f}") # noqa: E731 — "—" for null def ratio(v, sp, cp): return f"{v / sp:.1f}× / {v / cp:.1f}×" if (v is not None and sp > 0 and cp > 0) else "—" - cols = 4 + 2 * len(engines) # cls×2 + pipe×2 + one abs + one ×vs per engine - md = ["", "**Pre-tokenize: `classify + fsm` vs regex engines** — ns/byte, lower better. The fsm is " - "the scalar jump-table in both pipe columns; **SIMD / scalar is the classify pass** (regex " - "pre-tokenizers have no SIMD fsm). `×vs` = engine ÷ our pipeline (SIMD / scalar classify); " - "`onig` & `pcre2` (JIT) are C, `fancy` is pure-Rust fancy-regex, `logos` is a compile-time DFA " - "lexer (approximate grammar; n/a for deepseek).", "", - "| Fixture | classify SIMD | classify scalar | pipe (SIMD cls + fsm) | pipe (scalar cls + fsm) | " + cols = 4 + 2 * len(engines) # cls×2 + split×2 + one abs + one ×vs per engine + md = ["", "**Pre-tokenize: our split vs regex engines.** ns/byte, lower better. `split` is the " + "splitter this model's route actually runs, straight off the stage ladder. " + "`scalar-cls` swaps **only the classify pass** for its scalar version. It is not the split " + "without SIMD, since a regex-shaped fsm is a boundary-mask scanner and carries SIMD of its " + "own. `×vs` = engine ÷ our split (shipped / scalar-classify); `onig` & `pcre2` (JIT) are C, " + "`fancy` is pure-Rust fancy-regex, `logos` is a compile-time DFA lexer (approximate grammar; " + "n/a for deepseek).", "", + "| Fixture | classify SIMD | classify scalar | split (shipped) | split (scalar-cls) | " + " | ".join(engines) + " | " + " | ".join(f"×vs {e}" for e in engines) + " |", "|---" + "|---:" * cols + "|"] for r in sorted(rows, key=lambda r: (r["group"], r["fixture"])): diff --git a/tokenizers/tk-encode/examples/fixture_bench.rs b/tokenizers/tk-encode/examples/fixture_bench.rs index 9ac6d89bf3..947e566e5e 100644 --- a/tokenizers/tk-encode/examples/fixture_bench.rs +++ b/tokenizers/tk-encode/examples/fixture_bench.rs @@ -20,7 +20,10 @@ //! on, so the headline includes the post-process stage the ladder charges. //! 2. **Stage breakdown** — the `encode_generic::` ablation ladder plus //! the pre-tokenize-vs-regex-engine references, on fresh caller-owned -//! scratches, fully separate from the phase-1 timings. +//! scratches, fully separate from the phase-1 timings. Every rung drives the +//! same route a full `encode` takes (fused scan, zero-copy or staged, reported +//! as `route`), so the split bar is the splitter that actually ships and the +//! model bar is a difference between two runs of it. See `stage_secs`. //! 3. **Scaling & memory** — a multi-thread throughput sweep (1/2/4/8/max) over //! the whole corpus on fresh instances, and resident-set deltas measured by //! re-spawning this binary as `--memory ` children — one @@ -299,6 +302,23 @@ fn bench_throughput( /// successive levels and subtracting gives each stage's marginal cost (the ablation /// ladder), no profiler and no per-segment instrumentation. /// +/// Every rung runs the **route a full encode runs**, which is what makes subtracting +/// them meaningful. A fused pipeline's split rung drives the same masked FSM scan the +/// model rung does, emitting each span into `pre_tokens` instead of the model; a +/// zero-copy pipeline's split rung is its proven-cut pass, with neither the rewrite +/// write nor the added-token scan over rewritten text that the route never performs. +/// Two consequences to read the numbers with: +/// +/// - The split rung pays one span write per pre-token that a full fused encode does not +/// (it hands the span straight to the model), so the split bar reads a little high and +/// the model bar a little low. +/// - The `&str` slice the model call makes is charged to the model, not to the split. On +/// these corpora that is UTF-8 boundary checking, and it is not a small share. +/// +/// Before this was route-aware the split rung ran the *staged* array FSM, a splitter no +/// shipping encode uses. That put the split bar about twice too high and made the model +/// bar a subtraction across two different pipelines. +/// /// The scratch is created fresh here (never taken from the pipeline's pool), so the /// stage numbers are warmed on this fixture alone and can't perturb — or be flattered /// by — the phase-1 cache state. Both caller-owned buffers are reused across chunks @@ -333,6 +353,46 @@ fn stage_secs(pipeline: &PipelineTokenizer, chunks: &[String]) median_secs(samples) } +/// The top of the ladder with the output and span buffers allocated **per chunk**, the way +/// `PipelineTokenizer::encode` does for every call. The pipeline, the fresh scratch and +/// the warm-up are all identical to `stage_secs`, so subtracting `t_post` isolates +/// allocation with no other variable in play. +/// +/// It comes out near zero on these corpora: growing and freeing the output `Vec` is not a +/// cost worth chasing. It is reported anyway, because it looks like one. +/// +/// Timing the real `encode` here instead was tried and dropped. Phase 2 shares one pipeline +/// across every fixture, so its scratch pool holds the words of all previously benched +/// corpora, and against 64k slots that saturates. The delta then came out larger than the +/// whole encode: it measures cache eviction, not call overhead. A number that size in the +/// report would have read as a per-call cost. +fn fresh_buffer_secs(pipeline: &PipelineTokenizer, chunks: &[String]) -> f64 { + let mut scratch = pipeline.get_model().init_scratch(); + let mut run = || { + for chunk in chunks { + let mut out = Vec::new(); + let mut pre_tokens = Vec::new(); + let _ = pipeline.encode_generic::<{ PipelineTokenizer::STAGE_POSTPROCESS }>( + chunk, + true, + &mut pre_tokens, + &mut scratch, + &mut out, + ); + black_box(&out); + black_box(&pre_tokens); + } + }; + run(); // warm-up + let mut samples = Vec::with_capacity(REPS); + for _ in 0..REPS { + let start = Instant::now(); + run(); + samples.push(start.elapsed().as_secs_f64()); + } + median_secs(samples) +} + /// Stage decomposition + regex-engine references for one fixture: the /// `stage_ns_per_byte` and `pretok_vs_regex` objects of its report row. fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) -> (Value, Value) { @@ -341,10 +401,23 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) - let t_split = stage_secs::<{ PipelineTokenizer::STAGE_SPLIT }>(pipeline, &f.chunks); let t_model = stage_secs::<{ PipelineTokenizer::STAGE_MODEL }>(pipeline, &f.chunks); let t_post = stage_secs::<{ PipelineTokenizer::STAGE_POSTPROCESS }>(pipeline, &f.chunks); + let t_fresh = fresh_buffer_secs(pipeline, &f.chunks); // Two distinct "split" costs: `added_split` is the added/special-token scan (the // SpecialSegmentIterator over the AddedVocabulary, captured by the FRAME level), - // `pre_tokenize` is the pre-tokenizer split, `post` the special-token id-frame - // splice. All five stages sum exactly to `total`. + // `pre_tokenize` is the pre-tokenizer split *as this model's route performs it*, + // `post` the special-token id-frame splice. The five stages sum to `total`, which is + // the ladder's own full encode. + // + // `clamped` is what the five do NOT account for: the rungs are independent medians, so + // a stage at or below the noise floor goes negative and clamps to zero, after which the + // parts sum slightly above `total`. A `clamped` that is not near zero means this + // fixture's rungs are noise-dominated and its breakdown should not be trusted. + // + // `alloc` is the one cost outside the decomposition that can be isolated cleanly; see + // `fresh_buffer_secs`. `total` is still below what phase 1 measures through the public + // `encode`, which also acquires a pooled scratch and hands back an owned `Vec`; that + // remainder is NOT reported here, because no measurement in this bench separates it + // from the pool's cache state. Better an acknowledged gap than a mislabelled bar. let nspb = |secs: f64| secs * 1e9 / f.bytes as f64; let (ns_added, ns_norm, ns_split, ns_model, ns_post) = ( nspb(t_frame.max(0.0)), @@ -353,15 +426,23 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) - nspb((t_model - t_split).max(0.0)), nspb((t_post - t_model).max(0.0)), ); + let ns_total = nspb(t_post); + let clamped = ns_total - (ns_added + ns_norm + ns_split + ns_model + ns_post); + let ns_alloc = nspb(t_fresh - t_post); eprintln!( - " {} stages ns/byte: added-split {ns_added:.2}, norm {ns_norm:.2}, pre-split {ns_split:.2}, model {ns_model:.2}, post {ns_post:.2}", + " {} stages ns/byte: added-split {ns_added:.2}, norm {ns_norm:.2}, pre-split {ns_split:.2}, model {ns_model:.2}, post {ns_post:.2} → total {ns_total:.2} (clamped {clamped:+.2}) · buffer alloc {ns_alloc:+.2}", f.name ); - // pre_tokenize (= classify SIMD + fsm) vs classify-scalar and vs real regex engines - // over the same corpus, so the report shows the split beating a regex engine both - // WITH and WITHOUT SIMD. `scalar_pipe` = pre_tokenize + (cls_scalar − cls_simd): - // fsm is the scalar jump-table in both pipes, SIMD/scalar is the classify pass only. + // The shipped split vs real regex engines over the same corpus. `ns_split` is now the + // route's own splitter, so these ratios describe the code that runs. Before the ladder + // was route-aware they compared the engines against the staged array FSM and undersold + // the split by roughly 2x. + // + // `scalar_pipe` = ns_split + (cls_scalar − cls_simd) swaps *only* the classify pass + // for its scalar version. It is not "the split without SIMD": for a regex-shaped FSM + // the scan is a boundary-mask scanner, so SIMD lives in the scan too and cannot be + // subtracted out this way. Read it as the split with a scalar classify pass. let corpus: String = f.chunks.concat(); let cls_simd = classify_ns(corpus.as_bytes(), false); let cls_scalar = classify_ns(corpus.as_bytes(), true); @@ -384,7 +465,7 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) - }) }; eprintln!( - " {} pre-tok: SIMD-cls {ns_split:.2} / scalar-cls {scalar_pipe:.2} ns/B · vs onig {} · vs fancy {} · vs pcre2 {} · vs logos {}", + " {} split: {ns_split:.2} / scalar-classify {scalar_pipe:.2} ns/B · vs onig {} · vs fancy {} · vs pcre2 {} · vs logos {}", f.name, vs(onig_ns), vs(fancy_ns), @@ -400,7 +481,9 @@ fn bench_stages(pipeline: &PipelineTokenizer, f: &Fixture, regexes: &[String]) - "pre_tokenize": ns_split, "model": ns_model, "post": ns_post, - "total": nspb(t_post), + "total": ns_total, + "clamped": clamped, + "alloc": ns_alloc, }), json!({ "cls_simd": cls_simd, @@ -1188,6 +1271,9 @@ fn main() { models.push(json!({ "model": name, "desc": desc, "shape": shape, + // Which encode route the stage numbers decompose: the three split the work + // differently, so a stage chart is only readable against the route it timed. + "route": pipeline.encode_route().as_str(), "results": rows, "memory": memory, "threads": threads, "decode_threads": decode_threads, "decode_reason": decode_reason, })); diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index a43fc2048c..3a9a0e96da 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -300,6 +300,32 @@ pub(crate) enum FusedScan { DeepSeek, } +/// The route a chunk takes through [`PipelineTokenizer::encode_generic`]. See +/// [`PipelineTokenizer::encode_route`]. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum EncodeRoute { + /// The pending rewrite is read in place and proven cuts split the raw text: no + /// rewritten copy is built and no added-token scan runs over one. + ZeroCopy, + /// One native FSM cuts the chunk and hands each span straight to the BPE model, so + /// the split and the model are a single pass. + Fused, + /// Normalize, write, scan for added tokens, pre-tokenize into a span buffer, then run + /// the model once per span. + Staged, +} + +impl EncodeRoute { + /// The route's name, for a benchmark's report. + pub fn as_str(self) -> &'static str { + match self { + Self::ZeroCopy => "zero-copy", + Self::Fused => "fused", + Self::Staged => "staged", + } + } +} + impl PipelinePreTokenizer { /// The single native FSM this whole pre-tokenizer reduces to, if any: a recognized /// lone `Split`, or a `Sequence` that collapses to one. Each arm asks the child for @@ -887,6 +913,33 @@ impl PipelineTokenizer { &self.pre_tokenizer } + /// The single native FSM the whole split reduces to, when the model behind it is one + /// the fused scan can feed. [`encode_fused`](Self::encode_fused) and + /// [`encode_route`](Self::encode_route) both ask this, so the route a benchmark + /// reports cannot disagree with the route an encode takes. + fn fused_scan(&self) -> Option { + let scan = self.pre_tokenizer.native_fsm()?; + matches!(self.model, PipelineModel::BPE(_)).then_some(scan) + } + + /// Which route this tokenizer's chunks take through [`encode_generic`]. + /// + /// The three routes decompose into stages differently. The fused scan does the split + /// and the model in one pass, and the zero-copy route replaces the rewrite write and + /// the added-token scan over the rewritten text as well as the split. A stage + /// breakdown is only meaningful against the route it measured, so a benchmark reports + /// this next to the numbers. + pub fn encode_route(&self) -> EncodeRoute { + // Same order `encode_generic` tests them in. + if self.zero_copy.is_some() { + EncodeRoute::ZeroCopy + } else if self.fused_scan().is_some() { + EncodeRoute::Fused + } else { + EncodeRoute::Staged + } + } + /// Encode `input` into token ids. /// /// Special tokens are matched in two passes: @@ -991,21 +1044,32 @@ impl PipelineTokenizer { pending: None, } }; - // Full encodes consume a pending rewrite in place when the pipeline has - // the zero-copy route for it; the partial stages write it below, so the - // ablation ladder times the stages the zero-copy path replaces. - if STAGE >= Self::STAGE_POSTPROCESS - && normalized.pending.is_some() - && let Some(zero_copy) = &self.zero_copy - { - self.encode_chunk_zero_copy( - zero_copy, - &normalized.text, - pre_tokens, - scratch, - output, - )?; - continue; + // A pipeline with the zero-copy route reads the pending rewrite in + // place at *every* stage: it never writes the rewritten text and never + // scans it for added tokens, so no rung may charge this route for work + // it does not do. The cut is its split stage and the model its own. + // + // Whether the route applies is a property of the pipeline, not of + // the chunk. That matters below `STAGE_NORMALIZE`: there the ladder has + // not computed the rewrite yet, so `pending` is `None` for a reason + // that says nothing about the route. Falling through on a missing + // `pending` above that keeps ids right if a chunk ever lacks one. + if let Some(zero_copy) = &self.zero_copy { + if STAGE < Self::STAGE_NORMALIZE { + continue; + } + if normalized.pending.is_some() { + if STAGE >= Self::STAGE_SPLIT { + self.encode_chunk_zero_copy::( + zero_copy, + &normalized.text, + pre_tokens, + scratch, + output, + )?; + } + continue; + } } let normalized: Cow = normalized.write()?; @@ -1019,9 +1083,12 @@ impl PipelineTokenizer { } Segment::Text(normalized_chunk) => { if STAGE >= Self::STAGE_SPLIT { - if STAGE >= Self::STAGE_MODEL - && self.encode_fused(normalized_chunk, scratch, output) - { + if self.encode_fused::( + normalized_chunk, + pre_tokens, + scratch, + output, + ) { // Split and model ran as one pass; nothing left to do. } else { // Pre-tokenize the chunk of normalized text @@ -1059,13 +1126,23 @@ impl PipelineTokenizer { /// disappear, and the model reads each word while its bytes are still hot from /// the scan. Returns `false` when this tokenizer is not that shape, and the /// staged path must run instead. - fn encode_fused( + /// + /// Below [`STAGE_MODEL`](Self::STAGE_MODEL) the scan still runs, because it *is* the + /// split stage for a fused pipeline, and each span lands in `pre_tokens` instead of + /// going to the model. So the ladder's split rung times the scanner a full encode + /// really uses, and the model rung is a difference between two runs of the same scan. + /// Timing the staged array FSM here instead would report a splitter that no longer + /// runs, and make the model rung a subtraction across two different pipelines. + /// + /// [`STAGE_MODEL`]: Self::STAGE_MODEL + fn encode_fused( &self, chunk: &str, + pre_tokens: &mut Vec, scratch: &mut PipelineModelScratch, output: &mut Vec, ) -> bool { - let Some(scan) = self.pre_tokenizer.native_fsm() else { + let Some(scan) = self.fused_scan() else { return false; }; let PipelineModel::BPE(model) = &self.model else { @@ -1077,6 +1154,9 @@ impl PipelineTokenizer { thread_local! { static TAGS: RefCell> = const { RefCell::new(Vec::new()) }; } + if STAGE < Self::STAGE_MODEL { + pre_tokens.clear(); + } let bytes = chunk.as_bytes(); TAGS.with(|cell| { let tags = &mut *cell.borrow_mut(); @@ -1089,25 +1169,42 @@ impl PipelineTokenizer { scan_byte_level_masked, scan_cl100k_cap_masked, scan_deepseek_masked, scan_o200k_masked, scan_tekken_masked, }; + /// One span, either through the model or into the span buffer. `STAGE` is a + /// const generic, so a full encode compiles to the model call alone. + #[inline(always)] + fn feed( + model: &PipelineBPE, + chunk: &str, + span: Span, + pre_tokens: &mut Vec, + scratch: &mut BpeScratch, + output: &mut Vec, + ) { + if STAGE < PipelineTokenizer::STAGE_MODEL { + pre_tokens.push(span); + return; + } + model.tokenize_span(&chunk[span.range()], scratch, output); + } // One emit closure literal per arm: each scan gets its own instance by // value, which is what lets the emit inline into the scan loop. match scan { FusedScan::Gpt(GptFsm::Gpt2) => scan_byte_level_masked(bytes, tags, |span| { - model.tokenize_span(&chunk[span.range()], scratch, output); + feed::(model, chunk, span, pre_tokens, scratch, output); }), FusedScan::Gpt(GptFsm::Cl100k { digit_cap }) => { scan_cl100k_cap_masked(bytes, tags, digit_cap, |span| { - model.tokenize_span(&chunk[span.range()], scratch, output); + feed::(model, chunk, span, pre_tokens, scratch, output); }) } FusedScan::Gpt(GptFsm::O200k) => scan_o200k_masked(bytes, tags, |span| { - model.tokenize_span(&chunk[span.range()], scratch, output); + feed::(model, chunk, span, pre_tokens, scratch, output); }), FusedScan::Gpt(GptFsm::Tekken) => scan_tekken_masked(bytes, tags, |span| { - model.tokenize_span(&chunk[span.range()], scratch, output); + feed::(model, chunk, span, pre_tokens, scratch, output); }), FusedScan::DeepSeek => scan_deepseek_masked(bytes, tags, |span| { - model.tokenize_span(&chunk[span.range()], scratch, output); + feed::(model, chunk, span, pre_tokens, scratch, output); }), } }); @@ -1117,7 +1214,12 @@ impl PipelineTokenizer { /// One chunk through [`ZeroCopyMetaspace`]: cut the raw text, then read each span through /// the swap. The rewritten chunk is never built, and the special-token pass over it is /// skipped with it — the path is only built when no added token matches normalized text. - fn encode_chunk_zero_copy( + /// + /// The cut is this route's split stage, so below [`STAGE_MODEL`](Self::STAGE_MODEL) it + /// returns with `pre_tokens` filled and the model unrun. + /// + /// [`STAGE_MODEL`]: Self::STAGE_MODEL + fn encode_chunk_zero_copy( &self, zero_copy: &ZeroCopyMetaspace, chunk: &str, @@ -1131,6 +1233,9 @@ impl PipelineTokenizer { }; pre_tokens.clear(); let prepend = zero_copy.cut(chunk, pre_tokens); + if STAGE < Self::STAGE_MODEL { + return Ok(()); + } let mut spans = pre_tokens.iter(); if prepend && let Some(first) = spans.next() { // The prepend rewrites this one word for real, so its cache key is the rewritten @@ -1958,6 +2063,139 @@ mod tests { } } + /// The benchmark's stage ladder subtracts one rung from the next, which is only + /// meaningful if every rung runs the route a full encode runs. These three pin that + /// for each route, on the real configs. + /// + /// Without them the ladder can silently drift back to timing a splitter nothing uses: + /// the `pipeline_oracle` tests only ever drive `STAGE_POSTPROCESS`, so they prove ids + /// and say nothing about which path produced them. + mod stage_ladder { + use super::*; + + fn real_pipeline(file: &str) -> Option { + let path = format!("../data/{file}"); + if !std::path::Path::new(&path).exists() { + eprintln!("skip {file}: not present (fetch with `make bench-models`)"); + return None; + } + let tok = crate::Tokenizer::from_file(&path).unwrap(); + Some(PipelineTokenizer::try_from(&tok).unwrap()) + } + + const TEXT: &str = + "The quick brown fox jumps 123 times, don't it?\n\tif x == 1: return 'a'\n"; + + fn run(pipeline: &PipelineTokenizer) -> (Vec, Vec) { + let mut pre_tokens = Vec::new(); + let mut output = Vec::new(); + let mut scratch = pipeline.get_model().init_scratch(); + pipeline + .encode_generic::(TEXT, false, &mut pre_tokens, &mut scratch, &mut output) + .unwrap(); + (pre_tokens, output.iter().map(|t| t.id).collect()) + } + + /// The fused route's split rung must run the *fused scan* and cut the text exactly + /// where the staged pre-tokenizer does. + /// + /// Both halves are needed. Agreeing cuts alone would not catch the rung falling + /// back to the staged array FSM, since that FSM produces the same spans, and that + /// is exactly the bug this ladder had. So the fused entry point is called directly and + /// its `true` return asserted: `encode_generic` takes the staged path only when it + /// returns `false`, so a `true` here is proof of which splitter the rung timed. + #[test] + fn fused_split_rung_runs_the_fused_scan_and_cuts_where_the_pre_tokenizer_does() { + for file in ["gpt2.json", "llama-3-tokenizer.json"] { + let Some(pipeline) = real_pipeline(file) else { + continue; + }; + assert_eq!(pipeline.encode_route(), EncodeRoute::Fused, "{file}"); + + let mut fused_spans = Vec::new(); + let mut output = Vec::new(); + let mut scratch = pipeline.get_model().init_scratch(); + assert!( + pipeline.encode_fused::<{ PipelineTokenizer::STAGE_SPLIT }>( + TEXT, + &mut fused_spans, + &mut scratch, + &mut output, + ), + "{file}: the fused scan declined the split rung, so the ladder timed \ + the staged FSM instead" + ); + assert!( + output.is_empty(), + "{file}: the split rung must not run the model" + ); + + let mut staged = Vec::new(); + pipeline + .get_pre_tokenizer() + .pre_tokenize(TEXT, &mut staged) + .unwrap(); + assert_eq!(fused_spans, staged, "{file}"); + + // And the same spans come back out through the ladder itself. + let (through_ladder, ids) = run::<{ PipelineTokenizer::STAGE_SPLIT }>(&pipeline); + assert_eq!(through_ladder, staged, "{file}"); + assert!(ids.is_empty(), "{file}"); + } + } + + /// The zero-copy route's split rung is its proven-cut pass: spans out, no ids, + /// and no rewritten copy built. + #[test] + fn zero_copy_split_rung_cuts_without_running_the_model() { + for file in ["gemma-4.json", "llama-2.json"] { + let Some(pipeline) = real_pipeline(file) else { + continue; + }; + if pipeline.encode_route() != EncodeRoute::ZeroCopy { + continue; // `has_normalized_tokens` can veto the route for a config + } + let (spans, ids) = run::<{ PipelineTokenizer::STAGE_SPLIT }>(&pipeline); + assert!(!spans.is_empty(), "{file}: the cut produced nothing"); + assert!( + ids.is_empty(), + "{file}: the split rung must not run the model" + ); + } + } + + /// The top of the ladder is a real encode: `STAGE_MODEL` without specials has to + /// agree id-for-id with `encode`, on every route. This is what makes `total` + /// comparable to the phase-1 throughput number. + #[test] + fn model_rung_agrees_with_encode() { + for file in [ + "gpt2.json", + "llama-3-tokenizer.json", + "gemma-4.json", + "llama-2.json", + "bert-base-uncased.json", + ] { + let Some(pipeline) = real_pipeline(file) else { + continue; + }; + let (_, rung) = run::<{ PipelineTokenizer::STAGE_MODEL }>(&pipeline); + let direct: Vec = pipeline + .encode(TEXT, false) + .unwrap() + .iter() + .map(|t| t.id) + .collect(); + assert_eq!( + rung, + direct, + "{file} ({})", + pipeline.encode_route().as_str() + ); + } + } + } + #[test] fn segment_iterator_yields_text_and_specials_in_order() { let input = "aabbcc"; @@ -2381,4 +2619,3 @@ mod tests { assert_eq!(cache.lookup(b"hello").hit(), Some(&[7u32][..])); } } -