From d0b775082258edc11ff8095f24c83d7b9b0880eb Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 12:11:00 +0900 Subject: [PATCH 01/16] perf(bpe): hash a word once for both the fold and the cache probe `tokenize_pipeline` hashed every fold-missing pretoken twice: `fold_id` hashed it for the vocabulary's MPHF, then `WordCache::lookup` hashed the same bytes again for its home slot and tag. `BucketVocabStore` and `WordCache` seed `ahash` with the same four constants, so the second pass recomputed a value the first had already produced. Share it: `hash_word` exposes the value, `get_bytes_foldable_hashed` and `lookup_hashed` take it, and the model computes it once per word. Verification is untouched. The vocabulary still compares the entry's bytes to the query in full, and the cache still compares its key, so ids cannot change -- only the duplicated hash goes away. A word over fifteen bytes still pays the cache's second, independently seeded discriminant hash, which is what makes its key 127 bits rather than 64. The sharing is only sound while both sides seed identically, so a test pins them together. Re-seeding either would leave the cache placing a word under one hash and looking it up under another: no wrong ids, but every lookup would miss and the cache would quietly stop working. --- tokenizers/tk-encode/src/models/bpe/model.rs | 21 +++++++-- tokenizers/tk-encode/src/utils/word_cache.rs | 39 ++++++++++++++-- .../tk-encode/src/vocab/bucket_vocab_store.rs | 46 ++++++++++++++++++- 3 files changed, 99 insertions(+), 7 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 2f76ebfed..95c6d4805 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -244,9 +244,17 @@ impl PipelineBPE { /// entry that may be folded. `None` sends the word to the merge engines. #[inline(always)] fn fold_id(&self, sequence: &str) -> Option { + let bytes = sequence.as_bytes(); + self.fold_id_hashed(bytes, self.vocab.hash_word(bytes)) + } + + /// [`Self::fold_id`] for a caller that already hashed the word with + /// [`BucketVocabStore::hash_word`]. + #[inline(always)] + fn fold_id_hashed(&self, bytes: &[u8], hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable(sequence.as_bytes())?; + let (id, foldable) = self.vocab.get_bytes_foldable_hashed(bytes, hash)?; foldable.then_some(id) } @@ -309,7 +317,14 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - if let Some(id) = self.fold_id(sequence) { + // Hashed once for both probes. The fold asks the vocabulary and, on a miss, the cache asks + // its own table; `BucketVocabStore` and `WordCache` seed the same `ahash` state, so the two + // probes were hashing the same word to the same 64 bits twice. Both still verify their own + // way -- the vocabulary compares the entry's bytes, the cache compares its key -- so this + // shares the hash and nothing else. + let bytes = sequence.as_bytes(); + let hash = self.vocab.hash_word(bytes); + if let Some(id) = self.fold_id_hashed(bytes, hash) { output.push(PipelineToken { id }); return Ok(()); } @@ -322,7 +337,7 @@ impl pipeline::Model for PipelineBPE { // A word seen before costs a probe instead of a merge. let insert_at = if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { + match cache.lookup_hashed(bytes, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); return Ok(()); diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index a92ac3079..52e2ec462 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -122,11 +122,24 @@ impl<'a> WordCache { /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { + self.lookup_hashed(word, PLACEMENT_HASHER.hash_one(word)) + } + + /// [`Self::lookup`] for a caller that already hashed the word with [`placement_hash_of`]. + /// + /// Only the placement hash is handed in. The key still decides a hit, so a word of fifteen + /// bytes or fewer is compared to the slot exactly, as before. + pub fn lookup_hashed(&'a self, word: &[u8], placement_hash: u64) -> Lookup<'a> { + debug_assert_eq!( + placement_hash, + placement_hash_of(word), + "placement hash does not belong to this word" + ); let InsertPlacement { key, index: home, tag, - } = make_lookup_key(word, self.placement_mask); + } = make_lookup_key_hashed(word, placement_hash, self.placement_mask); let tag_window = self.tag_window(home); let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); @@ -364,9 +377,29 @@ impl<'a> SelfContained<'a> { #[repr(transparent)] pub struct LookupKey(u128); -/// The key, home slot and tag of a word. +/// The 64 bits a word's home slot and tag are taken from. +/// +/// Public because a caller that has to hash the same word for another table can compute this once +/// and hand it to [`WordCache::lookup_hashed`]. `BucketVocabStore` seeds its hasher identically, so +/// on the BPE path this is the value the vocabulary probe already produced. +#[inline] +pub fn placement_hash_of(word: &[u8]) -> u64 { + PLACEMENT_HASHER.hash_one(word) +} + +/// The key, home slot and tag of a word. Test-only: [`WordCache::lookup`] hashes the word itself +/// and goes straight to [`make_lookup_key_hashed`]. +#[cfg(test)] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { - let placement_hash = PLACEMENT_HASHER.hash_one(word); + make_lookup_key_hashed(word, placement_hash_of(word), placement_mask) +} + +/// [`make_lookup_key`] for a caller that already has the word's placement hash. +fn make_lookup_key_hashed( + word: &[u8], + placement_hash: u64, + placement_mask: u64, +) -> InsertPlacement { let key = if word.len() <= 15 { LookupKey::new_inline(word) } else { diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 7f62ec0ec..04a2de695 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -219,10 +219,28 @@ impl BucketVocabStore { /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { + self.get_bytes_foldable_hashed(q, self.hash_word(q)) + } + + /// The hash this store keys `q` by. Exposed so a caller that also has to hash the same word for + /// another table can pay for one pass instead of two; see [`Self::get_bytes_foldable_hashed`]. + #[inline] + pub fn hash_word(&self, q: &[u8]) -> u64 { + self.hasher.hash_one(q) + } + + /// [`Self::get_bytes_foldable`] for a caller that already hashed the word with + /// [`Self::hash_word`]. + /// + /// Verification is unchanged: the MPHF hands back a slot for *any* query, so the entry's bytes + /// are still compared to `q` in full. Only the hashing is shared, never the check. + #[inline] + pub fn get_bytes_foldable_hashed(&self, q: &[u8], hash: u64) -> Option<(u32, bool)> { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); + debug_assert_eq!(hash, self.hash_word(q), "hash does not belong to this word"); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; let (start, len) = (e.start as usize, e.len as usize); if len == q.len() && self.bytes[start..start + len] == *q { @@ -313,6 +331,32 @@ impl BucketVocabStore { mod tests { use super::*; + /// The BPE fast path hashes a word once and hands that one value to both the vocabulary probe + /// and the word cache, which is only sound while both seed `ahash` identically. If someone + /// re-seeds either side, the cache would place a word under one hash and look it up under + /// another: no wrong ids, but every lookup would miss and the cache would silently stop + /// working. This pins the two together so that change fails here instead. + #[test] + fn vocab_and_word_cache_hash_a_word_identically() { + let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]); + for w in [ + &b""[..], + b"a", + b"Hel", + b"the", + b" the", + b"fifteen bytes!!", + b"sixteen bytes ..", + b"a considerably longer word than the inline key can hold", + ] { + assert_eq!( + vocab.hash_word(w), + crate::utils::word_cache::placement_hash_of(w), + "hashers disagree on {w:?}" + ); + } + } + #[test] fn single_token() { let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]); From d64ad20b7dca1da0ef1c31c5f73d3d9629e2eb81 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:28:24 +0200 Subject: [PATCH 02/16] fix(bpe): repair `tokenize_spans`, which does not compile (#2314) #2304 added `PipelineBPE::tokenize_spans` against the model as it stood then. #2241 replaced the merge engines and #2310 dropped `ignore_merges`, and because the two landed on separate branches the merge produced a `feat/train_encode_split` that does not build: error[E0425]: cannot find type `Span` in this scope error[E0026]: struct `BpeScratch` does not have fields named `merge_queue`, `skip`, `word` error[E0027]: pattern does not mention fields `symbols`, `queue` error[E0609]: no field `ignore_merges` on type `&PipelineBPE` error[E0061]: this method takes 3 arguments but 4 arguments were supplied Bring the batch loop back in line with `tokenize_pipeline`: destructure `{ symbols, queue, word_cache }`, run the fold, and call the current `merge_word(sequence, symbols, queue)` followed by `unmap`. The fold has to stay ahead of the cache probe, as it is in `tokenize_pipeline`. A word that is a foldable vocabulary entry is answered in one probe and never enters the cache; probing the cache first would fill it with words the fold already serves for free, and the two paths would disagree about its contents. `tokenize_spans` overrides a trait method whose default is the `tokenize_pipeline` loop, so the two can drift without anything failing to build -- that is how this got in. Add a test that runs thousands of spans through one chunk (repeats, so the cache fills and hits; folded words; merged words; punctuation runs; multi-byte scripts; a long unbroken run) and compares the ids to the legacy reference. --- tokenizers/tk-encode/src/models/bpe/model.rs | 74 ++++++++++++++++---- 1 file changed, 60 insertions(+), 14 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4f5080c5c..3b1ec4709 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -8,7 +8,7 @@ use crate::models::bpe::legacy::model::BPE; use crate::models::bpe::merge_hot_cold_queue::{QueueScratch, merge_hot_cold_queue}; use crate::models::bpe::merge_multipass::merge_multipass; use crate::models::bpe::tables::BpeTables; -use crate::pipeline::{self, PipelineToken}; +use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; use crate::utils::word_cache::{Lookup, WordCache}; @@ -362,9 +362,8 @@ impl pipeline::Model for PipelineBPE { output: &mut Vec, ) -> Result<()> { let BpeScratch { - merge_queue, - skip, - word, + symbols, + queue, word_cache, } = scratch; @@ -380,6 +379,15 @@ impl pipeline::Model for PipelineBPE { continue; } + // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a + // word that is itself a foldable vocabulary entry in one probe, and those words never + // reach the cache. Probing the cache first would populate it with words the fold + // already serves for free, and the two paths would disagree about what it holds. + if let Some(id) = self.fold_id(sequence) { + output.push(PipelineToken { id }); + continue; + } + let mut placement = None; if let Some(cache) = word_cache.as_mut() { match cache.lookup(sequence.as_bytes()) { @@ -390,17 +398,13 @@ impl pipeline::Model for PipelineBPE { Lookup::Miss(at) => placement = Some(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`. + self.merge_word(sequence, symbols, queue); + // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids + output.extend(symbols.iter().map(|&symbol| PipelineToken { + id: self.tables.unmap.at(symbol as usize), + })); if let Some(cache) = word_cache.as_mut() && let Some(at) = placement { @@ -463,4 +467,46 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } + + /// `PipelineBPE::tokenize_spans` is an override of a trait method whose default is the + /// `tokenize_pipeline` loop, so the two can drift apart without anything failing to build -- + /// which is how it came to destructure a `BpeScratch` that no longer had those fields. + /// + /// The short strings above pass through the batch loop a handful of spans at a time. This one + /// gives it thousands in a single chunk, with the traffic that separates the two paths: + /// repeats (so the word cache both fills and hits), words the fold serves, words that must + /// merge, punctuation runs, multi-byte scripts, and a long unbroken run. + #[test] + fn the_batched_path_matches_the_reference() { + let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); + let pipe = PipelineTokenizer::try_from(&reference).unwrap(); + + let mut text = String::new(); + for i in 0..400 { + text.push_str(" the quick brown fox jumps over the lazy dog"); + text.push_str(" internationalisation unfortunately"); + text.push_str(" def foo(bar): return bar + 1"); + text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); + text.push_str(" 语言模型 ελληνικά"); + if i % 3 == 0 { + text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); + } + } + + let want: Vec = reference + .encode_fast(text.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipe + .encode(text.as_str(), false) + .wait() + .unwrap() + .remove(0) + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want.len(), got.len(), "token count differs"); + assert_eq!(want, got, "the batched path changed the ids"); + } } From 473db76fd0619ab947f5e1a650967ff9b658f852 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:17:22 +0200 Subject: [PATCH 03/16] perf(bpe): pack a short word into its key instead of hashing it (#2315) A pretoken is short. English averages 4.83 bytes of it, code 4.08, and the `<|...|>` shapes in `added-special-dense` 2.29. Running aHash over that is most of what the fold probe costs, and it buys nothing: the vocabulary compares the entry's bytes anyway, so the hash only has to spread well enough for the MPHF to separate keys. For a word of seven bytes or fewer, pack the bytes and the length into a `u64` and mix them with one multiply. Seven, so the length still fits in the top byte, which is what keeps `"ab"` from colliding with `"ab\0"`. Longer words keep aHash, which mixes the length in itself. `word_hash` is now the one definition. `BucketVocabStore::build`, every probe, and the word cache's placement all go through it, so a pretoken probed in both tables is hashed once for the pair and the two cannot drift apart. The per-struct `RandomState` goes away with it: consistency came from carrying the hasher around, and now it comes from there being a single function. Verification is unchanged and stays exact -- the vocabulary still compares the entry's bytes to the query in full, the cache still compares its 128-bit key. Nothing verifies with `mix`, which is why it does not have to be a strong hash. Dropping the mixing altogether does not work: packed short keys share their high bytes and MPHF construction fails with "indistinguishable hashes in bucket". Note this is why `WordCache::lookup` has to call `placement_hash_of` rather than a hasher of its own: the `debug_assert` in `lookup_hashed` caught exactly that mistake while this was being written. --- tokenizers/tk-encode/src/models/bpe/model.rs | 3 + tokenizers/tk-encode/src/utils/word_cache.rs | 25 ++--- .../tk-encode/src/vocab/bucket_vocab_store.rs | 96 ++++++++++++++++--- 3 files changed, 98 insertions(+), 26 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 95c6d4805..675cfc18f 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -242,6 +242,9 @@ impl PipelineBPE { /// The id to emit for `sequence` without merging, when the whole pretoken is a vocabulary /// entry that may be folded. `None` sends the word to the merge engines. + // Dead on this branch only: the batched `tokenize_spans` is its caller and this stack is based + // on a tip that predates it. Rebasing onto a base that has the batched path uses it again. + #[allow(dead_code)] #[inline(always)] fn fold_id(&self, sequence: &str) -> Option { let bytes = sequence.as_bytes(); diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 52e2ec462..335073818 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -56,17 +56,8 @@ use ahash::RandomState; use std::iter::Iterator; use wide::i8x16; -/// Hashes a word to the 64 bits its home slot and tag are taken from, and to the -/// bottom half of a long word's key ([`LookupKey::new_hash`]). -static PLACEMENT_HASHER: RandomState = RandomState::with_seeds( - 0x243f_6a88_85a3_08d3, - 0x1319_8a2e_0370_7344, - 0xa409_3822_299f_31d0, - 0x082e_fa98_ec4e_6c89, -); - /// Hashes a long word a second time, to fill the half of its key that -/// [`PLACEMENT_HASHER`] does not reach. The two hashes must be independent, or the +/// [`placement_hash_of`] does not reach. The two hashes must be independent, or the /// key would carry 64 bits of information instead of 127. static DISCRIMINANT_HASHER: RandomState = RandomState::with_seeds( 0x4528_21e6_38d0_1377, @@ -122,7 +113,7 @@ impl<'a> WordCache { /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { - self.lookup_hashed(word, PLACEMENT_HASHER.hash_one(word)) + self.lookup_hashed(word, placement_hash_of(word)) } /// [`Self::lookup`] for a caller that already hashed the word with [`placement_hash_of`]. @@ -379,12 +370,16 @@ pub struct LookupKey(u128); /// The 64 bits a word's home slot and tag are taken from. /// -/// Public because a caller that has to hash the same word for another table can compute this once -/// and hand it to [`WordCache::lookup_hashed`]. `BucketVocabStore` seeds its hasher identically, so -/// on the BPE path this is the value the vocabulary probe already produced. +/// This *is* [`crate::vocab::bucket_vocab_store::word_hash`], not a second function that happens to +/// agree with it. The BPE path probes the vocabulary and this cache for the same word, so it computes +/// the value once and hands it to [`WordCache::lookup_hashed`]; two definitions could drift and the +/// cache would place a word under one value and look it up under another. +/// +/// It also means a word of seven bytes or fewer is not hashed at all: the bytes and the length pack +/// into a `u64` and one multiply spreads them. #[inline] pub fn placement_hash_of(word: &[u8]) -> u64 { - PLACEMENT_HASHER.hash_one(word) + crate::vocab::bucket_vocab_store::word_hash(word) } /// The key, home slot and tag of a word. Test-only: [`WordCache::lookup`] hashes the word itself diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 04a2de695..fd21898d9 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -6,8 +6,8 @@ use std::fmt; type Mphf = FastPtrHash; -// Fixed seeds so a given vocab always hashes identically (the hasher is also stored on the struct, -// so build and query are guaranteed consistent regardless). +// Fixed seeds so a given vocab always hashes identically. Build and query both go through +// `word_hash`, which is the only place these are used, so they cannot drift apart. const SEEDS: [u64; 4] = [ 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, @@ -15,6 +15,85 @@ const SEEDS: [u64; 4] = [ 0x082E_FA98_EC4E_6C89, ]; +/// Hashes a word too long for the packed key. See [`word_hash`]. +static KEY_HASHER: RandomState = + RandomState::with_seeds(SEEDS[0], SEEDS[1], SEEDS[2], SEEDS[3]); + +/// How many bytes of a word fit in the packed key that replaces a hash pass. +/// +/// Seven, so the length still fits in the top byte of the same `u64`. That covers most pretokens: +/// english averages 4.83 bytes per pretoken, code 4.08, and the `<|...|>` shapes in +/// `added-special-dense` average 2.29. +const INLINE_KEY_BYTES: usize = 7; + +/// `KEY_MASK[len]` keeps the low `len` bytes of a `u64`; `LEN_TAG[len]` is the length in the top +/// byte. Tables rather than `u64::MAX >> (64 - 8 * len)` and `len << 56`: both loads are independent +/// of the word's bytes, so they overlap that load instead of queueing behind a shift chain. +const KEY_MASK: [u64; INLINE_KEY_BYTES + 1] = { + let mut m = [0u64; INLINE_KEY_BYTES + 1]; + let mut len = 1; + while len <= INLINE_KEY_BYTES { + m[len] = u64::MAX >> (64 - 8 * len); + len += 1; + } + m +}; +const LEN_TAG: [u64; INLINE_KEY_BYTES + 1] = { + let mut t = [0u64; INLINE_KEY_BYTES + 1]; + let mut len = 0; + while len <= INLINE_KEY_BYTES { + t[len] = (len as u64) << 56; + len += 1; + } + t +}; + +/// Mixes a packed short key into the well-distributed `u64` the MPHF and the cache's placement want. +/// +/// One multiply and one shift. This does not have to be a strong hash: nothing verifies with it. The +/// vocabulary still compares the entry's bytes and the cache still compares its key, so `mix` only +/// has to spread well enough to separate keys -- aHash's rounds are wasted on that. Dropping the +/// mixing entirely does *not* work: packed short keys share their high bytes, and MPHF construction +/// fails outright with "indistinguishable hashes in bucket". +#[inline(always)] +fn mix(z: u64) -> u64 { + let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); + z ^ (z >> 29) +} + +/// The `u64` a word is keyed by, in one pass and without a hash for a short word. +/// +/// Up to [`INLINE_KEY_BYTES`] bytes the word is packed into a `u64` with its length in the top byte +/// -- so `"ab"` cannot collide with `"ab\0"` -- and mixed with one multiply. Longer words fall back +/// to aHash, which mixes the length in itself. +/// +/// This is the single definition both `BucketVocabStore::build` and every probe go through, and +/// [`crate::utils::word_cache`] keys through it too, so a pretoken probed in both tables is hashed +/// once for the pair. It must stay one function: two copies could drift and the cache would place a +/// word under one value and look it up under another. +#[inline] +pub fn word_hash(word: &[u8]) -> u64 { + let len = word.len(); + if len > INLINE_KEY_BYTES { + return KEY_HASHER.hash_one(word); + } + // Reading past the word is not allowed, so read a head and a tail that overlap and stitch them: + // still register-only, no `memcpy`. + let raw = if len >= 4 { + let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; + let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u64; + head | tail << (8 * (len - 4)) + } else if len >= 1 { + let first = word[0] as u64; + let middle = (word[len / 2] as u64) << (8 * (len / 2)); + let last = (word[len - 1] as u64) << (8 * (len - 1)); + first | middle | last + } else { + 0 + }; + mix((raw & KEY_MASK[len]) | LEN_TAG[len]) +} + /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be /// emitted without running the merge loop. See `PipelineBPE::prove_fold`. /// @@ -55,7 +134,6 @@ struct Entry { #[derive(Clone)] pub struct BucketVocabStore { mphf: Mphf, - hasher: RandomState, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, /// `entries[slot]` -> (offset into `bytes`, length, id). Ordered by MPHF slot. @@ -102,12 +180,10 @@ impl BucketVocabStore { pub fn build(tokens: Vec<(Vec, u32)>) -> Self { let n = tokens.len(); - let hasher = RandomState::with_seeds(SEEDS[0], SEEDS[1], SEEDS[2], SEEDS[3]); - // 1. Pre-hash token bytes -> u64 keys using near perfect hash func let keys: Vec = tokens .iter() - .map(|(s, _)| hasher.hash_one(s.as_slice())) + .map(|(s, _)| word_hash(s.as_slice())) .collect(); // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely @@ -160,7 +236,7 @@ impl BucketVocabStore { *id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT" ); - let slot = mphf.index(&hasher.hash_one(s.as_slice())); + let slot = mphf.index(&word_hash(s.as_slice())); entries[slot] = Entry { start: bytes.len() as u32, len: s.len() as u16, @@ -172,7 +248,6 @@ impl BucketVocabStore { Self { mphf, - hasher, bytes: bytes.into_boxed_slice(), entries: entries.into_boxed_slice(), id_to_slot: id_to_slot.into_boxed_slice(), @@ -185,7 +260,6 @@ impl BucketVocabStore { let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), - hasher: RandomState::new(), bytes: Box::new([]), entries: Box::new([]), id_to_slot: Box::new([]), @@ -202,7 +276,7 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); + let slot = self.mphf.index(&word_hash(q)); let e = self.entries[slot]; let (start, len) = (e.start as usize, e.len as usize); @@ -226,7 +300,7 @@ impl BucketVocabStore { /// another table can pay for one pass instead of two; see [`Self::get_bytes_foldable_hashed`]. #[inline] pub fn hash_word(&self, q: &[u8]) -> u64 { - self.hasher.hash_one(q) + word_hash(q) } /// [`Self::get_bytes_foldable`] for a caller that already hashed the word with From 246d12291d8b1c2307022126bfc8b5b83e8b0c39 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:27:57 +0200 Subject: [PATCH 04/16] perf(bpe): put the key in the entry, so the fold probe is two loads (#2316) The fold probe was three dependent loads: the MPHF pilot, the entry, then the byte slab to compare the token against the query. The word cache does the same job in two, and the reason is layout, not luck -- its key lives in the slot it verifies, so nothing else has to be read. Give the vocabulary the same shape. `Entry` becomes `{ key, id }`, and `(start, len)` moves to a parallel `spans` array that only the reverse lookup and enumeration touch. A probe is now pilot + entry. Verification stays exact. A word of `INLINE_KEY_BYTES` or fewer has a key that *is* its bytes and its length, so comparing keys is proof of identity and the slab is never read. A longer word keys by aHash, which is not proof, so it still confirms against the slab -- the load it was paying anyway. So the saving lands exactly on the short pretokens that are the gap (english averages 4.83 bytes, code 4.08, `added-special-dense` 2.29) and nothing gives up the never-wrong guarantee. `LEN_TAG` now biases the length by one. A non-minimal MPHF returns padding slots, whose `Entry::default()` key is 0, and the probe rejects those with the same single compare it uses for everything else -- which only works while no real word can key to 0. The empty word keyed to exactly that before the bias. `key_and_hash` returns both halves so neither is recomputed: the model runs it once per word and hands the key and the hash to the fold probe and the hash to the cache. --- tokenizers/tk-encode/src/models/bpe/model.rs | 16 +- .../tk-encode/src/vocab/bucket_vocab_store.rs | 196 ++++++++++++++---- 2 files changed, 158 insertions(+), 54 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4421ecd21..b0ea136b4 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -12,7 +12,7 @@ use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; use crate::utils::word_cache::{Lookup, WordCache}; -use crate::vocab::bucket_vocab_store::BucketVocabStore; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -248,16 +248,16 @@ impl PipelineBPE { #[inline(always)] fn fold_id(&self, sequence: &str) -> Option { let bytes = sequence.as_bytes(); - self.fold_id_hashed(bytes, self.vocab.hash_word(bytes)) + let (key, hash) = key_and_hash(bytes); + self.fold_id_keyed(bytes, key, hash) } - /// [`Self::fold_id`] for a caller that already hashed the word with - /// [`BucketVocabStore::hash_word`]. + /// [`Self::fold_id`] for a caller that already ran [`key_and_hash`] on the word. #[inline(always)] - fn fold_id_hashed(&self, bytes: &[u8], hash: u64) -> Option { + fn fold_id_keyed(&self, bytes: &[u8], key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable_hashed(bytes, hash)?; + let (id, foldable) = self.vocab.get_bytes_foldable_keyed(bytes, key, hash)?; foldable.then_some(id) } @@ -326,8 +326,8 @@ impl pipeline::Model for PipelineBPE { // way -- the vocabulary compares the entry's bytes, the cache compares its key -- so this // shares the hash and nothing else. let bytes = sequence.as_bytes(); - let hash = self.vocab.hash_word(bytes); - if let Some(id) = self.fold_id_hashed(bytes, hash) { + let (key, hash) = key_and_hash(bytes); + if let Some(id) = self.fold_id_keyed(bytes, key, hash) { output.push(PipelineToken { id }); return Ok(()); } diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index fd21898d9..08e077874 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -42,7 +42,9 @@ const LEN_TAG: [u64; INLINE_KEY_BYTES + 1] = { let mut t = [0u64; INLINE_KEY_BYTES + 1]; let mut len = 0; while len <= INLINE_KEY_BYTES { - t[len] = (len as u64) << 56; + // `len + 1`, not `len`, so the top byte is never zero and no real key can be 0. A phantom + // slot's key is 0, which is what lets one compare reject it without a second load. + t[len] = (len as u64 + 1) << 56; len += 1; } t @@ -73,9 +75,20 @@ fn mix(z: u64) -> u64 { /// word under one value and look it up under another. #[inline] pub fn word_hash(word: &[u8]) -> u64 { + key_and_hash(word).1 +} + +/// The key a slot is verified by, and the `u64` the MPHF is indexed by. +/// +/// Up to [`INLINE_KEY_BYTES`] bytes the key *is* the word, so comparing it to a slot's key is +/// proof of identity and the byte slab never has to be read. Longer words key by aHash, which is +/// not proof, so those still confirm against the slab. +#[inline] +pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let len = word.len(); if len > INLINE_KEY_BYTES { - return KEY_HASHER.hash_one(word); + let hash = KEY_HASHER.hash_one(word); + return (hash, hash); } // Reading past the word is not allowed, so read a head and a tail that overlap and stitch them: // still register-only, no `memcpy`. @@ -91,7 +104,8 @@ pub fn word_hash(word: &[u8]) -> u64 { } else { 0 }; - mix((raw & KEY_MASK[len]) | LEN_TAG[len]) + let key = (raw & KEY_MASK[len]) | LEN_TAG[len]; + (key, mix(key)) } /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be @@ -105,14 +119,27 @@ const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -#[derive(Clone, Copy, Debug)] +/// Everything a probe reads, and nothing it does not. +/// +/// The key used to live in the byte slab, so verifying a slot meant a third dependent load after +/// the MPHF pilot and this entry. Holding it here makes the probe two loads. `(start, len)` moved +/// to [`Span`]: only the reverse lookup and enumeration want them, and keeping them here made +/// every probe drag six dead bytes through cache. +#[derive(Clone, Copy, Debug, Default)] struct Entry { - start: u32, - len: u16, + /// 0 for a phantom slot, which no real key can be -- see [`LEN_TAG`]. + key: u64, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } +/// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose. +#[derive(Clone, Copy, Debug, Default)] +struct Span { + start: u32, + len: u16, +} + /// The BucketVocabStore optimizes for space and speed. We don't use a HashMap to prevent duplicating the /// keys. Instead, we just use an `id_to_slot` and `entries` table. When you query bytes, you hash /// on the fly and get an `index` into the `entries` table. When you query an `id`, you fetch in @@ -136,8 +163,11 @@ pub struct BucketVocabStore { mphf: Mphf, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (offset into `bytes`, length, id). Ordered by MPHF slot. + /// `entries[slot]` -> (key, id). Ordered by MPHF slot. entries: Box<[Entry]>, + /// `spans[slot]` -> where the token's bytes live. Parallel to `entries`, read only by the + /// reverse lookup and by enumeration. + spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the @@ -218,14 +248,8 @@ impl BucketVocabStore { let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); - let mut entries = vec![ - Entry { - start: 0, - len: 0, - id: 0 - }; - n_slots - ]; + let mut entries = vec![Entry::default(); n_slots]; + let mut spans = vec![Span::default(); n_slots]; let mut id_to_slot = vec![u32::MAX; max_id as usize + 1]; for (s, id) in &tokens { assert!( @@ -236,11 +260,12 @@ impl BucketVocabStore { *id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT" ); - let slot = mphf.index(&word_hash(s.as_slice())); - entries[slot] = Entry { + let (key, hash) = key_and_hash(s.as_slice()); + let slot = mphf.index(&hash); + entries[slot] = Entry { key, id: *id }; + spans[slot] = Span { start: bytes.len() as u32, len: s.len() as u16, - id: *id, }; id_to_slot[*id as usize] = slot as u32; bytes.extend_from_slice(s); @@ -250,6 +275,7 @@ impl BucketVocabStore { mphf, bytes: bytes.into_boxed_slice(), entries: entries.into_boxed_slice(), + spans: spans.into_boxed_slice(), id_to_slot: id_to_slot.into_boxed_slice(), n, } @@ -262,6 +288,7 @@ impl BucketVocabStore { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), bytes: Box::new([]), entries: Box::new([]), + spans: Box::new([]), id_to_slot: Box::new([]), n: 0, } @@ -276,24 +303,38 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&word_hash(q)); - + let (key, hash) = key_and_hash(q); + let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - // Byte equality: confirms `q` really is the token at this slot (perfect hashing only - // guarantees a valid slot for in-vocab keys; this rejects collisions and Out Of Vocab queries). - if len == q.len() && self.bytes[start..start + len] == *q { - Some(e.id & VOCAB_ID_MASK) - } else { - None + // Perfect hashing only promises a valid slot for in-vocab keys, so the slot still has to be + // verified; this rejects collisions, phantom slots and out-of-vocabulary queries. + if e.key != key || !self.confirm(slot, q) { + return None; } + Some(e.id & VOCAB_ID_MASK) + } + + /// Whether the token at `slot` really is `q`. + /// + /// A word of [`INLINE_KEY_BYTES`] bytes or fewer has already proved it: its key *is* its bytes + /// and its length, so the caller's key compare was exact and this is free. Only a longer word, + /// whose key is a hash, reads the byte slab -- the load the probe used to pay unconditionally. + #[inline(always)] + fn confirm(&self, slot: usize, q: &[u8]) -> bool { + if q.len() <= INLINE_KEY_BYTES { + return true; + } + let s = self.spans[slot]; + let start = s.start as usize; + self.bytes.get(start..start + s.len as usize) == Some(q) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { - self.get_bytes_foldable_hashed(q, self.hash_word(q)) + let (key, hash) = key_and_hash(q); + self.get_bytes_foldable_keyed(q, key, hash) } /// The hash this store keys `q` by. Exposed so a caller that also has to hash the same word for @@ -303,25 +344,26 @@ impl BucketVocabStore { word_hash(q) } - /// [`Self::get_bytes_foldable`] for a caller that already hashed the word with - /// [`Self::hash_word`]. + /// [`Self::get_bytes_foldable`] for a caller that already ran [`key_and_hash`] on the word. /// /// Verification is unchanged: the MPHF hands back a slot for *any* query, so the entry's bytes /// are still compared to `q` in full. Only the hashing is shared, never the check. #[inline] - pub fn get_bytes_foldable_hashed(&self, q: &[u8], hash: u64) -> Option<(u32, bool)> { + pub fn get_bytes_foldable_keyed(&self, q: &[u8], key: u64, hash: u64) -> Option<(u32, bool)> { if self.entries.is_empty() { return None; } - debug_assert_eq!(hash, self.hash_word(q), "hash does not belong to this word"); + debug_assert_eq!( + (key, hash), + key_and_hash(q), + "key/hash pair does not belong to this word" + ); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - let (start, len) = (e.start as usize, e.len as usize); - if len == q.len() && self.bytes[start..start + len] == *q { - Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) - } else { - None + if e.key != key || !self.confirm(slot, q) { + return None; } + Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) } /// Records that this token folds to itself. Called once per entry at load, after the proof. @@ -345,9 +387,9 @@ impl BucketVocabStore { if slot == u32::MAX { return None; // id is within range but absent from the vocab } - let e = self.entries[slot as usize]; - let start = e.start as usize; - self.bytes.get(start..start + e.len as usize) + let s = self.spans[slot as usize]; + let start = s.start as usize; + self.bytes.get(start..start + s.len as usize) } #[inline] @@ -375,11 +417,12 @@ impl BucketVocabStore { } pub fn content(&self) -> Vec<(String, u32)> { - self.entries + self.spans .iter() - .filter(|e| e.len > 0) + .zip(self.entries.iter()) + .filter(|(s, _)| s.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(_, m)| m.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -391,11 +434,12 @@ impl BucketVocabStore { /// convenient when we want to re-build a vocab pub fn byte_content(&self) -> Vec<(Vec, u32)> { - self.entries + self.spans .iter() - .filter(|e| e.len > 0) + .zip(self.entries.iter()) + .filter(|(s, _)| s.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(_, m)| m.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() } @@ -431,6 +475,66 @@ mod tests { } } + /// A slot is verified by comparing its key. For a word of [`INLINE_KEY_BYTES`] bytes or fewer + /// the key *is* the bytes and the length, so that compare is proof and the byte slab is never + /// read; a longer word's key is a hash, so it still confirms against the slab. Both lengths + /// have to reject an out-of-vocabulary word, including ones that differ only in length. + #[test] + fn out_of_vocabulary_words_are_rejected() { + let words: [&[u8]; 7] = [ + b"a", + b"ab", + b"the", + b" the", + b"abcdefg", // exactly INLINE_KEY_BYTES: key is proof + b"abcdefgh", // one over: key is a hash, slab confirms + b"a much longer token than the inline key can hold", + ]; + let toks: Vec<(Vec, u32)> = words + .iter() + .enumerate() + .map(|(i, w)| (w.to_vec(), i as u32)) + .collect(); + let vocab = BucketVocabStore::build(toks.clone()); + + for (bytes, id) in &toks { + assert_eq!(vocab.get_bytes(bytes), Some(*id), "{bytes:?} should be found"); + } + for miss in [ + &b""[..], + b"b", + b"ba", + b"abc", + b"abcdef", // prefix of a present token + b"abcdefi", // same length as a present short token + b"abcdefghi",// same length class as a present long token + b"a much longer token than the inline key can hold!", + ] { + assert_eq!(vocab.get_bytes(miss), None, "{miss:?} is not in the vocab"); + } + } + + /// `Entry::default()` leaves key 0 in the padding slots a non-minimal MPHF returns, and the + /// probe rejects those with the same single compare it uses for everything else. That only + /// works while no real word can key to 0 -- which is why [`LEN_TAG`] biases the length by one. + #[test] + fn no_real_key_is_zero() { + for w in [ + &b""[..], + b"a", + b"\0", + b"\0\0\0\0\0\0\0", // seven zero bytes: only the length tag keeps this off 0 + b"1234567", + b"12345678", + ] { + assert_ne!( + key_and_hash(w).0, + 0, + "key of {w:?} collides with the phantom-slot sentinel" + ); + } + } + #[test] fn single_token() { let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]); From dcb313230edf02b628c9a5034d61626270fd6cf7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 13:32:38 +0900 Subject: [PATCH 05/16] perf(bpe): hash once on the batched path too `pipeline.rs` encodes through `tokenize_spans`, not `tokenize_pipeline`, and `tokenize_spans` was still hashing each word twice: `fold_id` for the vocabulary, then `cache.lookup` for the cache. Everything this PR does was landing only on `tokenize_pipeline`, which the encode loop does not call. Run `key_and_hash` once per word and hand the pair to `fold_id_keyed` and the hash to `lookup_hashed`, as `tokenize_pipeline` already does. `fold_id` had exactly one caller and folded into `fold_id_keyed` with it, taking a stale `#[allow(dead_code)]` with it. `the_batched_path_matches_the_reference` covers the path: ids compared against the legacy reference over thousands of spans in one chunk. --- tokenizers/tk-encode/src/models/bpe/model.rs | 24 ++++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index b0ea136b4..5882dc0ae 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -240,19 +240,11 @@ impl PipelineBPE { proven } - /// The id to emit for `sequence` without merging, when the whole pretoken is a vocabulary - /// entry that may be folded. `None` sends the word to the merge engines. - // Dead on this branch only: the batched `tokenize_spans` is its caller and this stack is based - // on a tip that predates it. Rebasing onto a base that has the batched path uses it again. - #[allow(dead_code)] - #[inline(always)] - fn fold_id(&self, sequence: &str) -> Option { - let bytes = sequence.as_bytes(); - let (key, hash) = key_and_hash(bytes); - self.fold_id_keyed(bytes, key, hash) - } - - /// [`Self::fold_id`] for a caller that already ran [`key_and_hash`] on the word. + /// The id to emit for a pretoken without merging, when the whole word is a vocabulary entry + /// that may be folded. `None` sends the word to the merge engines. + /// + /// Takes the key and hash rather than the word alone: both call sites also probe the cache for + /// the same bytes, so they run [`key_and_hash`] once and share it. #[inline(always)] fn fold_id_keyed(&self, bytes: &[u8], key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries @@ -401,14 +393,16 @@ impl pipeline::Model for PipelineBPE { // word that is itself a foldable vocabulary entry in one probe, and those words never // reach the cache. Probing the cache first would populate it with words the fold // already serves for free, and the two paths would disagree about what it holds. - if let Some(id) = self.fold_id(sequence) { + let bytes = sequence.as_bytes(); + let (key, hash) = key_and_hash(bytes); + if let Some(id) = self.fold_id_keyed(bytes, key, hash) { output.push(PipelineToken { id }); continue; } let mut placement = None; if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { + match cache.lookup_hashed(bytes, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); continue; From 574ee9bde7fd206ed6eee89b9f14abf1216b35b1 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:03:12 +0900 Subject: [PATCH 06/16] perf(pipeline): let the caller own the output buffer `encode_generic` sizes a fresh `Vec` from the input length -- a guess -- and hands back a new allocation on every call. A caller encoding many inputs (a batch, a server loop, a benchmark) can reserve once and `clear()` between calls instead: fewer allocations, and no first-touch of the token array each time. Split it: `encode_generic_into` takes `&mut Vec`, and `encode_generic` becomes the allocating wrapper, so nothing existing changes. Measured on identical code with both forms available, tokbench gpt2, 29 cells against gigatoken: 0.9233x allocating vs 0.9536x reusing -- ~3% of geomean throughput, ~6% on english (3.188 -> 3.001 ns/B). --- .../tk-encode/src/tokenizer/pipeline.rs | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 072efaf26..70b92719a 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -994,6 +994,24 @@ impl PipelineTokenizer { add_special_tokens: bool, ) -> Result> { let mut output = Vec::with_capacity(input.len() / 4); + self.encode_generic_into::(input, add_special_tokens, &mut output)?; + Ok(output) + } + + /// [`Self::encode_generic`] writing into a caller-owned buffer. + /// + /// The allocating form sizes its `Vec` from the input length, which is a guess, and hands back + /// a fresh allocation every call. A caller that encodes many inputs -- a batch, a server loop, + /// a benchmark -- can reserve once and `clear()` between calls instead: fewer allocations, and + /// no first-touch of the token array each time. Measured at ~3% of geomean throughput on + /// tokbench's 29 gpt2 cells (0.9233x -> 0.9536x against gigatoken, same code both sides). + #[doc(hidden)] + pub fn encode_generic_into( + &self, + input: &str, + add_special_tokens: bool, + output: &mut Vec, + ) -> Result<()> { let mut scratch = self.scratch_pool.get(&self.model); let PipelinePostProcessor { prefix, suffix } = &self.post_processor; // Prepend prefix tokens, if any @@ -1036,7 +1054,7 @@ impl PipelineTokenizer { normalized_chunk, &pre_tokens, &mut scratch, - &mut output, + output, )?; } Ok(()) @@ -1052,7 +1070,7 @@ impl PipelineTokenizer { if add_special_tokens && STAGE >= Self::STAGE_POSTPROCESS { output.extend_from_slice(suffix); } - Ok(output) + Ok(()) } } From 882af76791510a7ff47c36da5cadf45d7f00cfc5 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:07:35 +0900 Subject: [PATCH 07/16] merge: resolve duplicate tokenize_spans (keep the cache-aware, keyed version) The merge of #2313 into #2306 left two `tokenize_spans` definitions: git took both sides textually because they landed in different places. The stale one is #2306's, predating the word cache -- it destructures `BpeScratch { symbols, queue }` with no `word_cache` and calls the removed `fold_id`. Dropped it; kept the version that folds, probes the cache, and hashes each word once. --- tokenizers/tk-encode/src/models/bpe/model.rs | 37 -------------------- 1 file changed, 37 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 54901789f..5882dc0ae 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -426,43 +426,6 @@ impl pipeline::Model for PipelineBPE { Ok(()) } - /// Every pre-token of a chunk in one call. - /// - /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. - /// The scratch is destructured once instead of once per word, the output is grown once for the - /// whole batch instead of being capacity-checked on every push, and the virtual call, the - /// slice and the `Result` happen once per chunk rather than once per pre-token. - fn tokenize_spans( - &self, - chunk: &str, - spans: &[Span], - scratch: &mut Self::Scratch, - output: &mut Vec, - ) -> Result<()> { - let BpeScratch { symbols, queue } = scratch; - // One reservation for the batch. Most pre-tokens are a single token, so the span count is - // a close lower bound on what the batch emits; anything past it grows as usual. - output.reserve(spans.len()); - - for span in spans { - // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice - // of this chunk. Bounds- and UTF-8-checking it again per word measured worth removing. - let sequence = unsafe { chunk.get_unchecked(span.range()) }; - if sequence.is_empty() { - continue; - } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - continue; - } - self.merge_word(sequence, symbols, queue); - output.extend(symbols.iter().map(|&symbol| PipelineToken { - id: self.tables.unmap.at(symbol as usize), - })); - } - Ok(()) - } - fn init_scratch(&self) -> Self::Scratch { Self::Scratch { symbols: Vec::with_capacity(64), From be55698233db537afcde406c93a9586fec8133ab Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:28:43 +0900 Subject: [PATCH 08/16] perf(bpe): carry pair ranks across multipass passes, and stop splicing Two changes to the multipass engine, ported from the target-encode work. **Ranks carried across passes.** A pass used to re-look-up every pair it walked over, so the passes summed to O(n^2) table lookups -- measured at **41.9 per merged word**. Only the pairs touching a merge's product actually change, so `ranks[i]` (the value of the pair `(symbols[i], symbols[i+1])`) is now seeded by `convert_multipass` -- which already looks every pair up, so seeding costs one store per pair and no extra lookup -- and carried. A pass copies the ranks it did not invalidate and pays `get_value` **twice per merge** instead of once per symbol. Finding the next target is then a scan of `ranks` with no lookups at all. `prods` holds the matching product ids, kept apart so the search array stays a dense `u32` of ranks alone. **No memmove per merge.** A merge used to splice: write the product, then `copy_within` symbols, ranks and products to close the gap -- three memmoves on every merge. Instead the word carries a `live` bitmap of which slots still hold a symbol and a merge clears one bit; "previous live" and "next live" are `leading_zeros`/`trailing_zeros`. The `MAX_MP = 24` bound is what makes this work: it puts the live set in one `u64`. Dead pair slots hold `u32::MAX`, which is also "does not merge", so the minimum search skips them for free. The superseded sweep machinery (`MergeState`, `merge_once`, `batch_merging_is_safe`, `NOT_LEGAL`) goes with it -- the batching those implemented is subsumed by carrying ranks. Byte-exact: `the_proven_fold_never_changes_the_ids` and `the_batched_path_matches_the_reference` both compare ids against the legacy reference. --- .../tk-encode/src/models/bpe/convert.rs | 24 +- .../src/models/bpe/merge_multipass.rs | 314 ++++++++++-------- tokenizers/tk-encode/src/models/bpe/model.rs | 24 +- 3 files changed, 219 insertions(+), 143 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/convert.rs b/tokenizers/tk-encode/src/models/bpe/convert.rs index e47faa263..3a9235ca3 100644 --- a/tokenizers/tk-encode/src/models/bpe/convert.rs +++ b/tokenizers/tk-encode/src/models/bpe/convert.rs @@ -55,12 +55,22 @@ trait SinkMode { /// applies. struct MultipassSink<'a> { symbols: &'a mut Vec, + /// `ranks[i]` is the value of the pair `(symbols[i], symbols[i + 1])`. + /// + /// Conversion already looks every pair up, so seeding this costs one store per pair and no + /// extra lookup. It is what lets the merge passes stop re-ranking pairs they did not touch. + ranks: &'a mut Vec, + /// The product id of each pair, kept apart from its rank so the merge loop's search array stays + /// a dense `u32` of ranks alone. + prods: &'a mut Vec, lowest_merge: u64, } impl SinkMode for MultipassSink<'_> { #[inline(always)] fn record_pair(&mut self, merge: u64, _previous: u32, _symbol: u32) { + self.ranks.push((merge >> 32) as u32); + self.prods.push((merge & ID_MASK) as u32); self.lowest_merge = self.lowest_merge.min(merge); } #[inline(always)] @@ -129,13 +139,25 @@ impl SymbolSink { impl PipelineBPE { /// Converts one pretoken to internal IDs, returning the lowest-ranked adjacent pair, /// `u64::MAX` when no pair merges. - pub(super) fn convert_multipass(&self, sequence: &str, symbols: &mut Vec) -> u64 { + pub(super) fn convert_multipass( + &self, + sequence: &str, + symbols: &mut Vec, + ranks: &mut Vec, + prods: &mut Vec, + ) -> u64 { symbols.clear(); + ranks.clear(); + prods.clear(); // a word never has more symbols than bytes, so one reserve covers every push symbols.reserve(sequence.len()); + ranks.reserve(sequence.len()); + prods.reserve(sequence.len()); let mut sink = SymbolSink { mode: MultipassSink { symbols, + ranks, + prods, lowest_merge: u64::MAX, }, previous_symbol: u32::MAX, diff --git a/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs b/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs index b075d0e79..e9540423b 100644 --- a/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs +++ b/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs @@ -103,159 +103,199 @@ //! The merge is safe to batch merge only if the produced id (merged symbol) does not take part in other merges with lower rank (higher priority). //! This is enforced when building the lookup table and encoded in the `SAFE` bit. //! -use crate::models::bpe::tables::{BpeTables, ID_MASK, SAFE_MASK}; -use std::cmp; +use crate::models::bpe::tables::{BpeTables, ID_MASK}; -const NOT_LEGAL: u64 = u64::MAX; -/// Iteratively merges a word in place until it has no legal merge left -pub(super) fn merge_multipass(tables: &BpeTables, symbols: &mut Vec, mut target_merge: u64) { - let mut len = symbols.len(); - if len < 2 || target_merge == NOT_LEGAL { - return; - } - loop { - let MergeOnceOutput { - next_merge, - merged_length, - } = merge_once( - tables, - symbols, - len, - target_merge, - batch_merging_is_safe(tables, target_merge), - ); - len = merged_length; - if next_merge == NOT_LEGAL { - break; - } - target_merge = next_merge; +/// Iteratively merges a word in place until it has no legal merge left. +/// +/// `ranks[i]` is the value of the pair `(symbols[i], symbols[i + 1])`, seeded by +/// `convert_multipass` and carried across passes. That is the whole point: a pass used to re-look-up +/// every pair it walked over, so the passes summed to O(n^2) table lookups -- measured at 41.9 per +/// merged word. Only the pairs touching a merge's product actually change, so a pass now copies the +/// ranks it did not invalidate and pays `get_value` twice per merge instead of once per symbol. +/// Finding the next target is then a scan of `ranks`, with no lookups at all. +/// Wave merging is not an option here, and it is worth saying why: applying every *local* minimum in +/// one pass, rather than the one *global* minimum, is not byte-exact. Measured on 4 MB of english it +/// produced 1,473,346 tokens against the correct 944,838 -- a different tokenisation -- and was +/// slower besides (316 vs 525 MB/s). That closes chunk-wide bit-sliced rank comparison as a +/// direction: comparing all positions at once only helps if all the winners can be applied at once. +/// +/// 24 symbols: `GATE_ASCII` is 24 bytes and a word never has more symbols than bytes. +/// +/// Widening this to 64 (on a `u64` live mask) so the gate could move was tried, on the reasoning that +/// the gate existed because multipass searched a compacting array. It does not pay: routing long +/// words here instead of the hot/cold queue measured chinese 101 -> 93 MB/s and russian 152 -> 145 at +/// `GATE_MULTI = 64`. Multipass is O(n) search x O(n) merges against the queue's O(n log n), and by +/// n = 12 the queue is already ahead -- so the gate is not a leftover, it is the crossover. Keeping +/// the bound at 24 also keeps the three stack arrays to a 96-byte fill per word rather than 768. +const MAX_MP: usize = 24; + + +pub fn merge_multipass( + tables: &BpeTables, + symbols: &mut Vec, + ranks: &mut Vec, + prods: &mut Vec, + _first_merge: u64, +) { + let n = symbols.len(); + if n < 2 || n > MAX_MP { + return merge_multipass_vec(tables, symbols, ranks, prods); } - symbols.truncate(len); -} + debug_assert_eq!(ranks.len(), n - 1, "one rank per adjacent pair"); -/// Whether one pass may merge every occurrence of the target, or only the first occurrence. -#[inline(always)] -fn batch_merging_is_safe(tables: &BpeTables, target_merge: u64) -> bool { - // A NOT_LEGAL merge has the SAFE bit set, it would incorrectly return true here - // The caller is responsible for checking NOT_LEGAL does not reach this - debug_assert!(target_merge != NOT_LEGAL); - !tables.any_unsafe || (target_merge & SAFE_MASK != 0) -} + // Symbols keep their ORIGINAL positions for the whole merge; nothing is ever compacted. + // + // A merge used to splice: write the product, then `copy_within` symbols, ranks and products to + // close the gap -- three memmoves per merge, on every merge. Instead the word carries a `live` + // bitmap of which slots still hold a symbol, and a merge just clears one bit. Finding the symbol + // to the left or right of a position is then two bit ops rather than an index that shifted. + // + // The bound is what makes this work: it puts the live set in one `u64`, so "previous live" and + // "next live" are `leading_zeros` / `trailing_zeros`. Dead pair slots hold `u32::MAX`, which is + // also "does not merge", so the minimum search skips them for free and needs no mask of its own. + // The search runs to `n`, the word's real length -- padding it out to the bound was measured 3% + // slower, because the bound is 24 and the size is ~9. + // + // And it runs on the caller's buffers. Copying into `[u32; MAX_MP]` stack arrays first, which is + // what this did, compiled to a 288-byte `stp q` fill of all three arrays plus THREE out-of-line + // `memcpy` calls through the PLT per word, for at most 96 bytes each -- the bound only ever + // constrained `n`, never where the symbols had to live. + // Padding these out to `MAX_MP` so their length is a compile-time constant was tried, to let the + // bounds checks fold away and to hand the search a fixed-size reduction: 814 instructions instead + // of 640, and slower everywhere (english 3.40 -> 3.62 ns/B), because LLVM does not vectorise an + // argmin -- the index half of the reduction defeats it -- so the scan simply walked 23 padded + // lanes instead of 7. + let sym = &mut symbols[..]; + let rank = &mut ranks[..]; + let prod = &mut prods[..]; -/// One pass's cursors and running result. -/// -/// The read cursor marks the start of what is left of the old word, the write cursor -/// the end of the new word built so far (see the module docs). -struct MergeState { - read_cursor: usize, - write_cursor: usize, - target_merge: u64, - batched: bool, - was_merged: bool, - next_merge: u64, -} + // Bit i set means slot i still holds a symbol. `n <= 24`, so this never overflows. + let mut live: u64 = (1u64 << n) - 1; -impl MergeState { - fn new(target_merge: u64, batched: bool) -> Self { - Self { - read_cursor: 0, - write_cursor: 0, - target_merge, - batched, - was_merged: false, - next_merge: NOT_LEGAL, + #[inline(always)] + fn next_live(live: u64, from: usize) -> Option { + if from + 1 >= 64 { + return None; } + let above = live >> (from + 1); + (above != 0).then(|| from + 1 + above.trailing_zeros() as usize) } - - /// Looks up the pair at the read cursor and writes one symbol: the pair's product id when - /// its value equals the target and the pair may still merge, the left symbol otherwise. A - /// merge consumes both symbols of the pair, a copy only the left one. - /// - /// The returned value is the next call's `cached_pair_value`, and it is what keeps a step - /// down to a single table lookup. A copy has already paid for the lookup of - /// (`left_symbol`, `right_symbol`), and that is exactly the pair the next write has to - /// rank, so handing the value forward wins that second lookup back. A merge returns `None` - /// instead: it writes a product id that no pair has been looked up against yet, so the - /// next write has to pay for its own lookup. #[inline(always)] - fn step( - &mut self, - tables: &BpeTables, - symbols: &mut [u32], - cached_pair_value: Option, - ) -> Option { - let (left_symbol, right_symbol) = - (symbols[self.read_cursor], symbols[self.read_cursor + 1]); - let pair_value = tables.get_value(&left_symbol, &right_symbol); - let should_merge = pair_value == self.target_merge && (self.batched || !self.was_merged); - if should_merge { - self.was_merged = true; - self.read_cursor += 2; - let merged_symbol = (pair_value & ID_MASK) as u32; - self.write(tables, symbols, merged_symbol, None); - None - } else { - self.read_cursor += 1; - self.write(tables, symbols, left_symbol, cached_pair_value); - Some(pair_value) - } + fn prev_live(live: u64, before: usize) -> Option { + let below = live & ((1u64 << before) - 1); + (below != 0).then(|| 63 - below.leading_zeros() as usize) } - /// Writes one symbol at the write cursor and ranks the pair it forms with the previously - /// written symbol as a candidate for the next pass's target: `next_merge` keeps the lowest - /// value seen. `Some` reuses the value the caller already looked up, `None` pays for a - /// lookup here. The first written symbol has no left neighbour and nothing to rank. - #[inline(always)] - fn write( - &mut self, - tables: &BpeTables, - symbols: &mut [u32], - symbol: u32, - cached_pair_value: Option, - ) { - symbols[self.write_cursor] = symbol; - if self.write_cursor > 0 { - let rank = cached_pair_value - .unwrap_or_else(|| tables.get_value(&symbols[self.write_cursor - 1], &symbol)); - self.next_merge = cmp::min(self.next_merge, rank); + // The loop below indexes the caller's slices, so LLVM cannot see the bound and emits a + // bounds-check panic per access: 26 of them and 81 branches in a 640-instruction function. Every + // index is provably in range -- the scan gives `at < n - 1`, and `live` only ever has bits below + // `n` set -- so `get_unchecked` is sound here, and it takes the function to 555 instructions, 67 + // branches and 16 panics. It is still not worth the `unsafe`: measured 1.0031x over tokbench's 29 + // gpt2 cells, inside a +-0.8% noise floor. Whatever makes an iteration cost ~66 cycles -- for a + // scan over <= 8 `u32`, two bit ops, and two lookups that are demonstrably hot, since doubling + // `get_value` in place costs +2.1 ns of ~19 -- it is neither the bounds checks nor memory. + loop { + // Lowest rank, leftmost on a tie, over the word's real length. Dead slots are `u32::MAX`. + // + // This scan is already branchless: LLVM emits `cmp` + two `csel` (best, and the index), so + // the only branch is the perfectly-predicted loop-back. Packing the index into the low bits + // to turn the argmin into a plain min-reduction -- on the theory that the branch was the + // cost and that the index half was what blocked vectorisation -- measured **0.9836x** over + // tokbench's 29 gpt2 cells (english 3.26 -> 3.44), byte-exact. It backfired precisely + // because it succeeded: LLVM then vectorises the reduction, and the vector prologue (lane + // index vectors, four accumulators, a scalar epilogue) costs far more than the ~9 elements + // it processes. Do not retry either half of that idea without also pinning the trip count. + let mut best = u32::MAX; + let mut at = 0usize; + for (i, &r) in rank[..n - 1].iter().enumerate() { + if r < best { + best = r; + at = i; + } + } + if best == u32::MAX { + break; + } + + // The pair is (at, next_live(at)). The product lands in `at`; the right symbol dies. + let Some(right) = next_live(live, at) else { break }; + sym[at] = prod[at]; + live &= !(1u64 << right); + // No pair starts at the last symbol, so there is no rank slot for it; in the padded array + // this wrote `u32::MAX` over `u32::MAX`. + if right + 1 < n { + rank[right] = u32::MAX; // its pair is gone with it + } + rank[at] = u32::MAX; // recomputed below if a right neighbour remains + + // Only the pairs either side of the new symbol changed. + if let Some(pv) = prev_live(live, at) { + let v = tables.get_value(&sym[pv], &sym[at]); + rank[pv] = (v >> 32) as u32; + prod[pv] = (v & ID_MASK) as u32; + } + if let Some(nx) = next_live(live, at) { + let v = tables.get_value(&sym[at], &sym[nx]); + rank[at] = (v >> 32) as u32; + prod[at] = (v & ID_MASK) as u32; } - self.write_cursor += 1; } -} -struct MergeOnceOutput { - next_merge: u64, - merged_length: usize, + // Compact the survivors to the front, in position order. The write index never passes the read + // index, so this is safe in place and needs no second buffer. + let mut kept = 0usize; + let mut m = live; + while m != 0 { + let i = m.trailing_zeros() as usize; + sym[kept] = sym[i]; + kept += 1; + m &= m - 1; + } + symbols.truncate(kept); + ranks.clear(); + prods.clear(); } -/// One pass: walk the first `len` elements of `symbols` and merge occurrences of `target_merge`. -/// -/// Returns the next pass's target (`u64::MAX` when no pair in the rewritten word merges), and the number of symbols written. -/// Symbols past `len` are leftovers of earlier passes and should be truncated. -fn merge_once( +/// The compacting version, for the rare word longer than [`MAX_MP`] that still routes here. +fn merge_multipass_vec( tables: &BpeTables, - symbols: &mut [u32], - len: usize, - target_merge: u64, - batched: bool, -) -> MergeOnceOutput { - // Resliced so the loop bound and the slice length are the same value. Without this the - // compiler cannot connect `len` to the length of `symbols` and keeps a bounds check on - // every read and write of the sweep. - let symbols = &mut symbols[..len]; - let mut state = MergeState::new(target_merge, batched); - let mut cached_pair_value = None; - while state.read_cursor + 1 < len { - cached_pair_value = state.step(tables, symbols, cached_pair_value); - } - if state.read_cursor < len { - // The sweep's final symbol has no right neighbour to pair with, so it is copied as is. - let last_symbol = symbols[state.read_cursor]; - state.write(tables, symbols, last_symbol, cached_pair_value); - } - MergeOnceOutput { - next_merge: state.next_merge, - merged_length: state.write_cursor, + symbols: &mut Vec, + ranks: &mut Vec, + prods: &mut Vec, +) { + let mut len = symbols.len(); + while len >= 2 { + let mut best = u32::MAX; + let mut at = 0usize; + for (i, &rank) in ranks[..len - 1].iter().enumerate() { + if rank < best { + best = rank; + at = i; + } + } + if best == u32::MAX { + break; + } + symbols[at] = prods[at]; + symbols.copy_within(at + 2..len, at + 1); + len -= 1; + if len >= 2 { + ranks.copy_within(at + 1..len, at); + prods.copy_within(at + 1..len, at); + } + if at > 0 { + let v = tables.get_value(&symbols[at - 1], &symbols[at]); + ranks[at - 1] = (v >> 32) as u32; + prods[at - 1] = (v & ID_MASK) as u32; + } + if at + 1 < len { + let v = tables.get_value(&symbols[at], &symbols[at + 1]); + ranks[at] = (v >> 32) as u32; + prods[at] = (v & ID_MASK) as u32; + } } + symbols.truncate(len); + ranks.truncate(len.saturating_sub(1)); + prods.truncate(len.saturating_sub(1)); } diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 5882dc0ae..4f98172e9 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -219,6 +219,8 @@ impl PipelineBPE { let mut proven = vec![false; len]; let mut symbols = Vec::with_capacity(64); let mut scratch = QueueScratch::default(); + let mut ranks = Vec::with_capacity(64); + let mut prods = Vec::with_capacity(64); for id in 0..len as u32 { let Some(bytes) = self.vocab.id_to_token_bytes(id) else { continue; @@ -232,7 +234,7 @@ impl PipelineBPE { // A single atom has no pair to merge and is trivially its own encoding. true } else { - self.merge_word(text, &mut symbols, &mut scratch); + self.merge_word(text, &mut symbols, &mut scratch, &mut ranks, &mut prods); symbols.len() == 1 && self.tables.unmap.at(symbols[0] as usize) == id }; proven[id as usize] = foldable; @@ -264,6 +266,8 @@ impl PipelineBPE { sequence: &str, symbols: &mut Vec, queue_scratch: &mut QueueScratch, + ranks: &mut Vec, + prods: &mut Vec, ) { let bytes = sequence.as_bytes(); // Classify on the first content byte, not on the delimiter the pre-tokenizer prepended. @@ -279,8 +283,8 @@ impl PipelineBPE { ); merge_hot_cold_queue(&self.tables, symbols, queue_scratch); } else { - let first_merge = self.convert_multipass(sequence, symbols); - merge_multipass(&self.tables, symbols, first_merge); + let first_merge = self.convert_multipass(sequence, symbols, ranks, prods); + merge_multipass(&self.tables, symbols, ranks, prods, first_merge); } } } @@ -292,6 +296,10 @@ pub struct BpeScratch { pub(crate) symbols: Vec, /// Entry arena and the two queue tiers. pub(crate) queue: QueueScratch, + /// One rank per adjacent pair of the word being merged; see `merge_multipass`. + pub(crate) ranks: Vec, + /// The matching product ids, kept apart so the merge loop's search array is ranks alone. + pub(crate) prods: Vec, /// Words already seen, so a repeat costs a probe instead of a merge. It lives in the scratch /// so it outlives the encode call that fills it -- otherwise it would never see a word twice. pub(crate) word_cache: Option, @@ -327,6 +335,8 @@ impl pipeline::Model for PipelineBPE { let BpeScratch { symbols, queue, + ranks, + prods, word_cache, } = scratch; @@ -344,7 +354,7 @@ impl pipeline::Model for PipelineBPE { }; let start = output.len(); - self.merge_word(sequence, symbols, queue); + self.merge_word(sequence, symbols, queue, ranks, prods); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), @@ -374,6 +384,8 @@ impl pipeline::Model for PipelineBPE { let BpeScratch { symbols, queue, + ranks, + prods, word_cache, } = scratch; @@ -412,7 +424,7 @@ impl pipeline::Model for PipelineBPE { } let start = output.len(); - self.merge_word(sequence, symbols, queue); + self.merge_word(sequence, symbols, queue, ranks, prods); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), @@ -430,6 +442,8 @@ impl pipeline::Model for PipelineBPE { Self::Scratch { symbols: Vec::with_capacity(64), queue: QueueScratch::default(), + ranks: Vec::with_capacity(64), + prods: Vec::with_capacity(64), word_cache: self.cache_capacity.map(WordCache::new), } } From 3508f8339eaab1e97d8147f4c31e3ba72665bfdf Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:34:29 +0900 Subject: [PATCH 09/16] perf(bpe): write cache hits straight at the output cursor A cache hit went through the tag row -- one load of 16 control bytes, a SIMD compare, then the slot -- and handed back a slice the caller walked. Both are avoidable for the common case, a word cached in its own home slot with at most three ids. `probe_emit_hashed` reads the home slot directly and stores all `MAX_INLINE_IDS` lanes unconditionally at a `*mut u32` the caller supplies, so the line is touched once and the ids never become a slice. Lanes past `ids_len` are dead: the caller advances its cursor by `ids_len` only, so the next word overwrites them or the final `set_len` cuts them off. A spilled or off-home slot falls back to the window walk without re-keying the word (`lookup_placed`, split out of `lookup_hashed`). `tokenize_spans` keeps a raw cursor and one capacity check per word covering both the fold's single write and the probe's lanes, and reserves `2 * spans.len() + MAX_INLINE_IDS` up front -- 92% of english pre-tokens are one id and 98% at most two, so the old `spans.len()` was a lower bound that made the buffer grow, and memcpy what it held, partway through most chunks. Deliberately NOT taken from the source branch: its `LookupKey` is a `u64`, which makes a hit on a word over seven bytes a 2^-64 proposition. This keeps the 128-bit key, so a hit stays exact for words up to fifteen bytes as before; only the emit is fused. Byte-exact: `the_batched_path_matches_the_reference` drives thousands of spans through this path, cache hits included, and compares ids to the legacy reference under debug assertions. --- tokenizers/tk-encode/src/models/bpe/model.rs | 59 ++++++++++++--- tokenizers/tk-encode/src/utils/word_cache.rs | 75 +++++++++++++++++++- 2 files changed, 124 insertions(+), 10 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4f98172e9..c36f3faf1 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -11,7 +11,7 @@ use crate::models::bpe::tables::BpeTables; use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; -use crate::utils::word_cache::{Lookup, WordCache}; +use crate::utils::word_cache::{Lookup, MAX_INLINE_IDS, ProbeEmit, WordCache}; use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash}; const GATE_MULTI: u16 = 8; @@ -389,9 +389,12 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // One reservation for the batch. Most pre-tokens are a single token, so the span count is - // a close lower bound on what the batch emits; anything past it grows as usual. - output.reserve(spans.len()); + // 92% of english pre-tokens are one id and 98% at most two, so reserve for two apiece plus + // the probe's headroom. `spans.len()` alone is a *lower* bound, which would make the buffer + // grow -- and memcpy what it already holds -- partway through most chunks. + output.reserve(2 * spans.len() + MAX_INLINE_IDS); + let mut capacity = output.capacity(); + let mut cursor = output.len(); for span in spans { // SAFETY: the pre-tokenizer cuts on char boundaries, so a span is always a valid slice @@ -401,6 +404,15 @@ impl pipeline::Model for PipelineBPE { continue; } + // One capacity check per word, covering both the fold's single write and the probe's + // `MAX_INLINE_IDS` lanes. After it, writing that many past `cursor` is in bounds. + if cursor + MAX_INLINE_IDS > capacity { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; + output.reserve(spans.len() + MAX_INLINE_IDS); + capacity = output.capacity(); + } + // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a // word that is itself a foldable vocabulary entry in one probe, and those words never // reach the cache. Probing the cache first would populate it with words the fold @@ -408,21 +420,48 @@ impl pipeline::Model for PipelineBPE { let bytes = sequence.as_bytes(); let (key, hash) = key_and_hash(bytes); if let Some(id) = self.fold_id_keyed(bytes, key, hash) { - output.push(PipelineToken { id }); + // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. + unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; + cursor += 1; continue; } let mut placement = None; if let Some(cache) = word_cache.as_mut() { - match cache.lookup_hashed(bytes, hash) { - Lookup::Hit(ids) => { + // The probe writes the ids at the cursor itself, so a hit is one load of the slot + // and one unconditional store of its lanes -- the ids never become a slice and the + // line is never read twice. + // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and + // `PipelineToken` is layout-identical to `u32` (asserted at the top of this file). + let found = unsafe { + cache.probe_emit_hashed( + bytes, + hash, + output.as_mut_ptr().add(cursor).cast::(), + ) + }; + match found { + ProbeEmit::Wrote(n) => { + cursor += n; + continue; + } + // A hit the fast path could not serve: the probe already found the ids, so copy + // those rather than probing a second time. + ProbeEmit::Hit(ids) => { + // SAFETY: `cursor` counts what has been written so far. + unsafe { output.set_len(cursor) }; output.extend(ids.iter().map(|&id| PipelineToken { id })); + cursor = output.len(); + capacity = output.capacity(); continue; } - Lookup::Miss(at) => placement = Some(at), + ProbeEmit::Miss(at) => placement = Some(at), } } + // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` + // through its normal API, so its length has to be true again first. + unsafe { output.set_len(cursor) }; let start = output.len(); self.merge_word(sequence, symbols, queue, ranks, prods); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -434,7 +473,11 @@ impl pipeline::Model for PipelineBPE { { cache.insert(at, output[start..].iter().map(|token| token.id)); } + cursor = output.len(); + capacity = output.capacity(); } + // SAFETY: `cursor` counts every token written, by the fast paths and the slow one alike. + unsafe { output.set_len(cursor) }; Ok(()) } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 335073818..798753618 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -126,11 +126,66 @@ impl<'a> WordCache { placement_hash_of(word), "placement hash does not belong to this word" ); + self.lookup_placed(make_lookup_key_hashed(word, placement_hash, self.placement_mask)) + } + + /// A hit written straight to `dst`, skipping the tag window. + /// + /// The common case is a word cached in its own home slot with at most [`MAX_INLINE_IDS`] ids. + /// `lookup` reaches that through the tag row -- one load of 16 control bytes, a SIMD compare, + /// then the slot itself -- and hands back a slice the caller walks. Both are avoidable: read + /// the home slot directly and store all [`MAX_INLINE_IDS`] lanes unconditionally, so the line + /// is touched once and the ids never become a slice. + /// + /// Lanes past `ids_len` are dead. The caller advances its cursor by `ids_len` only, so the next + /// word overwrites them, or the final `set_len` cuts them off. + /// + /// # Safety + /// + /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. + #[inline] + pub unsafe fn probe_emit_hashed( + &'a self, + word: &[u8], + placement_hash: u64, + dst: *mut u32, + ) -> ProbeEmit<'a> { + debug_assert_eq!( + placement_hash, + placement_hash_of(word), + "placement hash does not belong to this word" + ); + let placement = make_lookup_key_hashed(word, placement_hash, self.placement_mask); + // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`) and the table is + // `next_pow2 + WINDOW_SIZE` long, so the home slot is always in bounds. + let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; + // An untouched slot holds `LookupKey(0)`, which no word keys to, so a key match here is a + // real hit -- the same argument the window walk makes. A spilled slot is excluded because + // its payload holds offsets, not ids, and only spilled slots can be stale. + if slot.key == placement.key && !slot.is_spilled() { + // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. + unsafe { + for lane in 0..MAX_INLINE_IDS { + dst.add(lane).write(slot.payload[lane]); + } + } + return ProbeEmit::Wrote(slot.ids_len as usize); + } + match self.lookup_placed(placement) { + Lookup::Hit(ids) => ProbeEmit::Hit(ids), + Lookup::Miss(at) => ProbeEmit::Miss(at), + } + } + + /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup_hashed`] + /// so [`Self::probe_emit_hashed`] can fall back to it without re-keying the word. + #[inline] + fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { key, index: home, tag, - } = make_lookup_key_hashed(word, placement_hash, self.placement_mask); + } = placement; let tag_window = self.tag_window(home); let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); @@ -167,7 +222,7 @@ impl<'a> WordCache { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; - let word = if len <= 3 { + let word = if len <= MAX_INLINE_IDS { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { @@ -484,6 +539,22 @@ pub struct InsertPlacement { tag: u8, } +/// How many ids a slot holds inline. A slot is 32 bytes: a `u128` key, three `u32` ids and a +/// length, so three is what fits beside the key. +pub const MAX_INLINE_IDS: usize = 3; + +/// What [`WordCache::probe_emit_hashed`] found. `Wrote` is the fast path: the ids are already at +/// `dst` and only the count comes back. +pub enum ProbeEmit<'a> { + /// An inline hit in the home slot. [`MAX_INLINE_IDS`] lanes were written at `dst`; this many + /// of them are live. + Wrote(usize), + /// A hit the fast path could not serve -- a spilled entry, or one placed off its home slot. + /// The ids were found, so the caller copies these rather than probing again. + Hit(&'a [u32]), + Miss(InsertPlacement), +} + pub enum Lookup<'a> { Hit(&'a [u32]), Miss(InsertPlacement), From 61c82c1dc870b1361f32225cfff7003ee8b6fc20 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:41:08 +0900 Subject: [PATCH 10/16] perf(bpe): take the target-encode word cache and vocabulary store u64 packed keys throughout, digest verification in the vocabulary store, and the pipelined probe helpers (probe_slot/entry_at/resolve_foldable). Accepts the exactness trades deliberately: a cache hit on a word over seven bytes is 2^-64, and an out-of-vocabulary pretoken can be mistaken for a vocabulary token at 2^-32, where both were previously impossible. --- tokenizers/tk-encode/src/models/bpe/model.rs | 17 +- .../tk-encode/src/pre_tokenizers/split.rs | 12 +- tokenizers/tk-encode/src/utils/word_cache.rs | 267 +++++------- .../tk-encode/src/vocab/bucket_vocab_store.rs | 411 ++++++++---------- 4 files changed, 308 insertions(+), 399 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index c36f3faf1..2f5524a20 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -248,10 +248,11 @@ impl PipelineBPE { /// Takes the key and hash rather than the word alone: both call sites also probe the cache for /// the same bytes, so they run [`key_and_hash`] once and share it. #[inline(always)] - fn fold_id_keyed(&self, bytes: &[u8], key: u64, hash: u64) -> Option { + fn fold_id_keyed(&self, key: u64, hash: u64) -> Option { // One probe; the foldable bit is part of the id that probe already returned. Which entries - // carry it was settled at load -- see `from_bpe`. - let (id, foldable) = self.vocab.get_bytes_foldable_keyed(bytes, key, hash)?; + // carry it was settled at load -- see `from_bpe`. The key verifies the slot, so the bytes + // are not needed here at all. + let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } @@ -327,7 +328,7 @@ impl pipeline::Model for PipelineBPE { // shares the hash and nothing else. let bytes = sequence.as_bytes(); let (key, hash) = key_and_hash(bytes); - if let Some(id) = self.fold_id_keyed(bytes, key, hash) { + if let Some(id) = self.fold_id_keyed(key, hash) { output.push(PipelineToken { id }); return Ok(()); } @@ -342,7 +343,7 @@ impl pipeline::Model for PipelineBPE { // A word seen before costs a probe instead of a merge. let insert_at = if let Some(cache) = word_cache.as_mut() { - match cache.lookup_hashed(bytes, hash) { + match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); return Ok(()); @@ -419,7 +420,7 @@ impl pipeline::Model for PipelineBPE { // already serves for free, and the two paths would disagree about what it holds. let bytes = sequence.as_bytes(); let (key, hash) = key_and_hash(bytes); - if let Some(id) = self.fold_id_keyed(bytes, key, hash) { + if let Some(id) = self.fold_id_keyed(key, hash) { // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; cursor += 1; @@ -434,8 +435,8 @@ impl pipeline::Model for PipelineBPE { // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and // `PipelineToken` is layout-identical to `u32` (asserted at the top of this file). let found = unsafe { - cache.probe_emit_hashed( - bytes, + cache.probe_emit_keyed( + key, hash, output.as_mut_ptr().add(cursor).cast::(), ) diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 04a6b4c30..740aa984d 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -193,6 +193,11 @@ impl pipeline::PreTokenizer for Split { // everything else keeps the FSM. match fsm { GptFsm::Gpt2 if bitsplit::fast_builder() => { + // SCRATCH -- NOT FOR COMMIT: prove which splitter actually runs. + if std::env::var_os("TK_ROUTE").is_some() { + static ONCE: std::sync::Once = std::sync::Once::new(); + ONCE.call_once(|| eprintln!("[route] gpt2 -> bitsplit (SIMD)")); + } pipeline::classify_into_spans_bits( text.as_bytes(), bitsplit::bitsplit_byte_level, @@ -208,7 +213,12 @@ impl pipeline::PreTokenizer for Split { ); return Ok(()); } - _ => {} + _ => { + if std::env::var_os("TK_ROUTE").is_some() { + static ONCE2: std::sync::Once = std::sync::Once::new(); + ONCE2.call_once(|| eprintln!("[route] fell through to the FSM")); + } + } } pipeline::classify_into_spans( text.as_bytes(), diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 798753618..b8c2982cc 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -52,19 +52,23 @@ use std::fmt::Debug; -use ahash::RandomState; use std::iter::Iterator; use wide::i8x16; -/// Hashes a long word a second time, to fill the half of its key that -/// [`placement_hash_of`] does not reach. The two hashes must be independent, or the -/// key would carry 64 bits of information instead of 127. -static DISCRIMINANT_HASHER: RandomState = RandomState::with_seeds( - 0x4528_21e6_38d0_1377, - 0xbe54_66cf_34e9_0c6c, - 0xc0ac_29b7_c97c_50dd, - 0x3f84_d5b5_b547_0917, -); +use crate::vocab::bucket_vocab_store::key_and_hash; +#[cfg(test)] +use crate::vocab::bucket_vocab_store::INLINE_KEY_BYTES; + +// One hash pass, not two, and shared with the vocabulary. +// +// This used to hash a long word twice -- `PLACEMENT_HASHER` for its slot and `DISCRIMINANT_HASHER` +// for the other half of its key -- which is why long pretokens paid for the key. It now keys through +// `bucket_vocab_store::key_and_hash`, the same function the fold probes with, so a word that misses +// the fold and falls through to this table is hashed once for both rather than once each. + +/// How many ids a [`WordCacheSlot`] holds inline before it has to spill. A probe writes this +/// many lanes unconditionally, so it is also the headroom [`WordCache::probe_emit`] needs. +pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) pub struct WordCache { @@ -112,58 +116,60 @@ impl<'a> WordCache { /// Looks up a word in the cache. /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted + #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { - self.lookup_hashed(word, placement_hash_of(word)) + self.lookup_placed(make_lookup_key(word, self.placement_mask)) } - /// [`Self::lookup`] for a caller that already hashed the word with [`placement_hash_of`]. - /// - /// Only the placement hash is handed in. The key still decides a hit, so a word of fifteen - /// bytes or fewer is compared to the slot exactly, as before. - pub fn lookup_hashed(&'a self, word: &[u8], placement_hash: u64) -> Lookup<'a> { - debug_assert_eq!( - placement_hash, - placement_hash_of(word), - "placement hash does not belong to this word" - ); - self.lookup_placed(make_lookup_key_hashed(word, placement_hash, self.placement_mask)) + /// [`Self::lookup`] for a caller that already has the word's key and hash from + /// [`key_and_hash`] -- the fold probes the vocabulary with the same pair, so sharing it means a + /// word that misses the fold and falls through to here is hashed once, not twice. + #[inline] + pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { + self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) } - /// A hit written straight to `dst`, skipping the tag window. + /// Probe and emit in one step: on an inline hit in the home slot the ids are written straight + /// to `dst` and the count returned, so nothing goes back to the slot and nothing becomes a + /// slice. /// - /// The common case is a word cached in its own home slot with at most [`MAX_INLINE_IDS`] ids. - /// `lookup` reaches that through the tag row -- one load of 16 control bytes, a SIMD compare, - /// then the slot itself -- and hands back a slice the caller walks. Both are avoidable: read - /// the home slot directly and store all [`MAX_INLINE_IDS`] lanes unconditionally, so the line - /// is touched once and the ids never become a slice. + /// This is the shape the hot path wants. [`Self::lookup`] hands back a `&[u32]`, which means + /// the caller re-reads the slot to build a fat pointer and then copies a run whose length it + /// only learns at run time -- three trips over one 32-byte line that a single load already + /// brought in. Here that line is read once, all [`MAX_INLINE_IDS`] lanes are stored + /// unconditionally, and the caller advances its cursor by the count: no branch on the length, + /// no second load, no slice. /// - /// Lanes past `ids_len` are dead. The caller advances its cursor by `ids_len` only, so the next - /// word overwrites them, or the final `set_len` cuts them off. + /// Falls back to the full window walk for anything else. The table is sized well above its + /// load, so a word's home slot is usually the one it was placed in and the walk is a few + /// percent of words. /// /// # Safety + /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. `word` must not be empty -- + /// an empty word keys to zero, which is also what an untouched slot holds. + #[inline] + pub unsafe fn probe_emit(&'a self, word: &[u8], dst: *mut u32) -> ProbeEmit<'a> { + debug_assert!(!word.is_empty(), "probe_emit needs a non-empty word"); + let (key, hash) = key_and_hash(word); + unsafe { self.probe_emit_keyed(key, hash, dst) } + } + + /// [`Self::probe_emit`] for a caller that already has the word's key and hash. /// - /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. + /// # Safety + /// As [`Self::probe_emit`]: `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. #[inline] - pub unsafe fn probe_emit_hashed( - &'a self, - word: &[u8], - placement_hash: u64, - dst: *mut u32, - ) -> ProbeEmit<'a> { - debug_assert_eq!( - placement_hash, - placement_hash_of(word), - "placement hash does not belong to this word" - ); - let placement = make_lookup_key_hashed(word, placement_hash, self.placement_mask); - // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`) and the table is + pub unsafe fn probe_emit_keyed(&'a self, key: u64, hash: u64, dst: *mut u32) -> ProbeEmit<'a> { + let placement = placement_from(LookupKey(key), hash, self.placement_mask); + // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is // `next_pow2 + WINDOW_SIZE` long, so the home slot is always in bounds. let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; - // An untouched slot holds `LookupKey(0)`, which no word keys to, so a key match here is a - // real hit -- the same argument the window walk makes. A spilled slot is excluded because - // its payload holds offsets, not ids, and only spilled slots can be stale. + // An untouched slot holds `LookupKey(0)`, which no non-empty word can key to, so a key + // match here is a real hit -- the same 127-bit argument the window walk makes. if slot.key == placement.key && !slot.is_spilled() { - // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. + // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. Lanes past `ids_len` are + // dead: the caller advances its cursor by `ids_len` only, so the next word overwrites + // them or the final `set_len` cuts them off. unsafe { for lane in 0..MAX_INLINE_IDS { dst.add(lane).write(slot.payload[lane]); @@ -177,8 +183,8 @@ impl<'a> WordCache { } } - /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup_hashed`] - /// so [`Self::probe_emit_hashed`] can fall back to it without re-keying the word. + /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup`] so + /// [`Self::probe_emit`] can fall back to it without hashing the word a second time. #[inline] fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { @@ -323,7 +329,7 @@ impl WordCacheSlot { } pub fn new_self_contained(key: LookupKey, ids: impl ExactSizeIterator) -> Self { - assert!(ids.len() <= 3); + assert!(ids.len() <= MAX_INLINE_IDS); let ids_len = ids.len() as u8; let mut payload = [0u32; 3]; for (slot, id) in payload.iter_mut().zip(ids) { @@ -421,115 +427,42 @@ impl<'a> SelfContained<'a> { /// ``` #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] #[repr(transparent)] -pub struct LookupKey(u128); +pub struct LookupKey(u64); -/// The 64 bits a word's home slot and tag are taken from. +/// The key, home slot and tag of a word. /// -/// This *is* [`crate::vocab::bucket_vocab_store::word_hash`], not a second function that happens to -/// agree with it. The BPE path probes the vocabulary and this cache for the same word, so it computes -/// the value once and hands it to [`WordCache::lookup_hashed`]; two definitions could drift and the -/// cache would place a word under one value and look it up under another. -/// -/// It also means a word of seven bytes or fewer is not hashed at all: the bytes and the length pack -/// into a `u64` and one multiply spreads them. +/// The key and its hash come from [`key_and_hash`], which is also what the vocabulary store keys on +/// -- one scheme, so a word that is probed in both tables can be hashed once. See +/// [`WordCache::lookup_keyed`]. #[inline] -pub fn placement_hash_of(word: &[u8]) -> u64 { - crate::vocab::bucket_vocab_store::word_hash(word) -} - -/// The key, home slot and tag of a word. Test-only: [`WordCache::lookup`] hashes the word itself -/// and goes straight to [`make_lookup_key_hashed`]. -#[cfg(test)] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { - make_lookup_key_hashed(word, placement_hash_of(word), placement_mask) + let (key, hash) = key_and_hash(word); + placement_from(LookupKey(key), hash, placement_mask) } -/// [`make_lookup_key`] for a caller that already has the word's placement hash. -fn make_lookup_key_hashed( - word: &[u8], - placement_hash: u64, - placement_mask: u64, -) -> InsertPlacement { - let key = if word.len() <= 15 { - LookupKey::new_inline(word) - } else { - LookupKey::new_hash(DISCRIMINANT_HASHER.hash_one(word), placement_hash) - }; +#[inline] +fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlacement { InsertPlacement { key, - index: (placement_hash & placement_mask) as usize, - tag: ((placement_hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), + index: (hash & placement_mask) as usize, + tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } -impl LookupKey { - pub const TAG_MASK: u128 = 1 << 127; - - /// The key of a word of fifteen bytes or fewer: the word is its own key. - pub fn new_inline(word: &[u8]) -> Self { - let len = word.len(); - assert!(len <= 15); - // yes, this is a bit weird :) - // - // We used to do this: - // ```rust - // payload[..word.len()].copy_from_slice(word); - // payload[15] = word.len() as u8; - // Self(u128::from_le_bytes(payload)) - // ``` - // But that would compile into a memcpy call, probably because the len is only known at runtime. - // memcpy turned out to be quite slow and inefficient. - // - // The head / tail with fixed size compiles into plain register loads which are way faster - let raw = 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 << (8 * (len - 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 << (8 * (len - 4)) - } else if len >= 1 { - let first = word[0] as u128; - let middle = (word[len / 2] as u128) << (8 * (len / 2)); - let last = (word[len - 1] as u128) << (8 * (len - 1)); - first | middle | last - } else { - 0 - }; - Self(raw | (len as u128) << 120) - } - - /// The key of a longer word: 127 bits of hash stand in for the word's bytes. - pub fn new_hash(discriminant: u64, placement: u64) -> Self { - Self(Self::TAG_MASK | (discriminant as u128) << 64 | placement as u128) - } -} - impl std::fmt::Debug for LookupKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut debug = f.debug_struct("LookupKey"); - if self.0 & Self::TAG_MASK == 0 { + // An inline key carries its length in the top byte and its bytes below; anything else is a + // hash and has nothing readable in it. + let len = (self.0 >> 56) as usize; + if len <= 7 { let bytes = self.0.to_le_bytes(); - let len = bytes[15] as usize; - debug - .field("type", &"inline") - .field("word_len", &len) - .field("word", &bytes[..len].escape_ascii().to_string()); + f.debug_tuple("LookupKey") + .field(&String::from_utf8_lossy(&bytes[..len]).into_owned()) + .finish() } else { - debug - .field("type", &"hashed") - .field( - "discriminant", - &format!("{:#x}", ((self.0 & !Self::TAG_MASK) >> 64) as u64), - ) - .field( - "placement", - &format!("{:#x}", (self.0 & u64::MAX as u128) as u64), - ); + write!(f, "LookupKey(hash {:#018x})", self.0) } - debug.finish() } } @@ -539,12 +472,13 @@ pub struct InsertPlacement { tag: u8, } -/// How many ids a slot holds inline. A slot is 32 bytes: a `u128` key, three `u32` ids and a -/// length, so three is what fits beside the key. -pub const MAX_INLINE_IDS: usize = 3; +pub enum Lookup<'a> { + Hit(&'a [u32]), + Miss(InsertPlacement), +} -/// What [`WordCache::probe_emit_hashed`] found. `Wrote` is the fast path: the ids are already at -/// `dst` and only the count comes back. +/// What [`WordCache::probe_emit`] found. `Wrote` is the fast path: the ids are already at the +/// caller's cursor and only the count comes back. pub enum ProbeEmit<'a> { /// An inline hit in the home slot. [`MAX_INLINE_IDS`] lanes were written at `dst`; this many /// of them are live. @@ -555,11 +489,6 @@ pub enum ProbeEmit<'a> { Miss(InsertPlacement), } -pub enum Lookup<'a> { - Hit(&'a [u32]), - Miss(InsertPlacement), -} - struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, @@ -733,9 +662,19 @@ mod tests { fn packed_keys_are_unique_per_word() { let cache = WordCache::new(1 << 8); let key = |word: &[u8]| make_lookup_key(word, cache.placement_mask).key; + // Past the inline range these are hashes, so distinctness is the hash's job, not the + // packing's -- but they must still not collide. assert_ne!(key(b"aaaaaaaaaaaaaa\x7f"), key(b"aaaaaaaaaaaaaa\xff")); + // Inside it, the length is part of the key, so a trailing NUL cannot be lost. assert_ne!(key(b"abcd"), key(b"abcd\0")); - assert_eq!(key(b"aaaaaaaaaaaaaa\xff").0 & LookupKey::TAG_MASK, 0); + // Every distinct word within the inline range gets a distinct key, by construction. + let mut seen = std::collections::HashSet::new(); + for len in 1..=INLINE_KEY_BYTES { + for b in 0..=255u8 { + let word: Vec = (0..len).map(|i| b.wrapping_add(i as u8)).collect(); + assert!(seen.insert(key(&word).0), "collision at len={len} b={b}"); + } + } } /// One window shape per row: the needle in various lanes, an empty slot in @@ -815,19 +754,21 @@ mod tests { } /// Every inline length, with a different value in every byte position, so a - /// packing that drops, duplicates or misplaces a byte fails. The reference is - /// the construction the packing must be equivalent to: the bytes copied into - /// a zeroed array, the length written in the top byte. + /// packing that drops, duplicates or misplaces a byte fails. The reference is the construction + /// the packing must be equivalent to: the bytes in a zeroed `u64`, the length in the top byte. + /// + /// Only up to [`INLINE_KEY_BYTES`]; past that the key is a hash and there are no bytes in it to + /// check. #[test] fn an_inline_key_is_the_words_bytes_with_the_length_on_top() { - for len in 0..=15usize { + for len in 0..=INLINE_KEY_BYTES { let word: Vec = (1..=len as u8).collect(); - let mut padded = [0u8; 16]; + let mut padded = [0u8; 8]; padded[..len].copy_from_slice(&word); - padded[15] = len as u8; + padded[7] = len as u8; assert_eq!( - LookupKey::new_inline(&word), - LookupKey(u128::from_le_bytes(padded)), + key_and_hash(&word).0, + u64::from_le_bytes(padded), "len={len}" ); } @@ -866,7 +807,7 @@ mod tests { assert_ne!(tag, WordCache::EMPTY, "pick a word with a nonzero tag"); cache.quick_lookup[index] = tag; cache.cached_words[index] = - WordCacheSlot::new_self_contained(LookupKey::new_inline(b"decoy"), [7].into_iter()); + WordCacheSlot::new_self_contained(LookupKey(key_and_hash(b"decoy").0), [7].into_iter()); assert_eq!(cache.lookup(b"beta").hit(), None); store(&mut cache, b"beta", &[2]); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 08e077874..896c7ab34 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,85 +4,94 @@ use ahash::RandomState; use ptr_hash::{FastPtrHash, PtrHashParams, hash::NoHash}; use std::fmt; -type Mphf = FastPtrHash; - -// Fixed seeds so a given vocab always hashes identically. Build and query both go through -// `word_hash`, which is the only place these are used, so they cannot drift apart. -const SEEDS: [u64; 4] = [ +/// Hashes a word key. Fixed seeds so a vocabulary always hashes identically. +/// +/// One pass. xxh3-128 was tried, to give a long key 63 bits of discrimination independent of the 64 +/// that place it; it cost 5.8 -> 7.0 ns per probe and ~10-25% on chinese and russian, whose pretokens +/// are mostly long, so it was dropped. A long key reuses its placement hash as its discriminant and +/// adds the length instead: a false hit needs a 64-bit collision at equal length (~2^-64 per query) +/// rather than being impossible as the old `memcmp` made it. +static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, 0xA409_3822_299F_31D0, 0x082E_FA98_EC4E_6C89, -]; +); + +type Mphf = FastPtrHash; -/// Hashes a word too long for the packed key. See [`word_hash`]. -static KEY_HASHER: RandomState = - RandomState::with_seeds(SEEDS[0], SEEDS[1], SEEDS[2], SEEDS[3]); +// No hasher on the struct: both hashes below are fixed, so build and query agree without one +// having to be carried along to keep them consistent. -/// How many bytes of a word fit in the packed key that replaces a hash pass. +/// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be +/// emitted without running the merge loop. See `PipelineBPE::prove_fold`. /// -/// Seven, so the length still fits in the top byte of the same `u64`. That covers most pretokens: -/// english averages 4.83 bytes per pretoken, code 4.08, and the `<|...|>` shapes in -/// `added-special-dense` average 2.29. -const INLINE_KEY_BYTES: usize = 7; - -/// `KEY_MASK[len]` keeps the low `len` bytes of a `u64`; `LEN_TAG[len]` is the length in the top -/// byte. Tables rather than `u64::MAX >> (64 - 8 * len)` and `len << 56`: both loads are independent -/// of the word's bytes, so they overlap that load instead of queueing behind a shift chain. -const KEY_MASK: [u64; INLINE_KEY_BYTES + 1] = { - let mut m = [0u64; INLINE_KEY_BYTES + 1]; - let mut len = 1; - while len <= INLINE_KEY_BYTES { - m[len] = u64::MAX >> (64 - 8 * len); - len += 1; - } - m -}; -const LEN_TAG: [u64; INLINE_KEY_BYTES + 1] = { - let mut t = [0u64; INLINE_KEY_BYTES + 1]; - let mut len = 0; - while len <= INLINE_KEY_BYTES { - // `len + 1`, not `len`, so the top byte is never zero and no real key can be 0. A phantom - // slot's key is 0, which is what lets one compare reject it without a second load. - t[len] = (len as u64 + 1) << 56; - len += 1; - } - t -}; +/// Packed into the id rather than kept beside it, following the same convention as a packed merge +/// value, where the low bits are the product id and the high bits are a flag field (`SAFE_MASK`). +/// The lookup already loads this entry to verify the key, so the flag costs no memory, no extra +/// load, and no second structure to keep in step with the vocabulary. +const FOLD_BIT: u32 = 1 << 31; +/// The id half. 2^31 ids is far past any vocabulary. +const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -/// Mixes a packed short key into the well-distributed `u64` the MPHF and the cache's placement want. +/// Tokens up to this many bytes are their own key: the bytes fit beside the length in a `u64`. /// -/// One multiply and one shift. This does not have to be a strong hash: nothing verifies with it. The -/// vocabulary still compares the entry's bytes and the cache still compares its key, so `mix` only -/// has to spread well enough to separate keys -- aHash's rounds are wasted on that. Dropping the -/// mixing entirely does *not* work: packed short keys share their high bytes, and MPHF construction -/// fails outright with "indistinguishable hashes in bucket". +/// Seven, not fifteen, because that is what the corpus is. English averages 4.83 bytes per +/// pretoken and code 4.08, so a `u128` key was paying double width, a two-part head/tail read and a +/// 32-byte entry to describe words that fit in a register. +pub(crate) const INLINE_KEY_BYTES: usize = 7; + +/// Mixes a short key into the well-distributed `u64` the MPHF wants. +/// +/// One multiply and one shift. An inline key *is* the token, so the compare is exact no matter how +/// the slot was chosen -- the hash only has to spread well enough for the MPHF to separate the keys, +/// and aHash's rounds, or splitmix64's second multiply, are wasted on that. Dropping the mixing +/// entirely does not work: packed short keys share their high bytes, and construction fails outright +/// with "indistinguishable hashes in bucket". #[inline(always)] fn mix(z: u64) -> u64 { let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); z ^ (z >> 29) } -/// The `u64` a word is keyed by, in one pass and without a hash for a short word. +/// The token's bytes as a fixed-width key, and the `u64` the MPHF is indexed by, from one pass. +/// +/// * up to [`INLINE_KEY_BYTES`] bytes: the key *is* the bytes, with the length in the top byte, so +/// the compare is proof and `"ab"` cannot collide with `"ab\0"`. +/// * longer: the key is aHash of the bytes, which mixes the length in, and doubles as the placement +/// hash. A false hit then needs a full 64-bit collision (~2^-64 per query) rather than being +/// impossible as a `memcmp` against the byte arena made it. /// -/// Up to [`INLINE_KEY_BYTES`] bytes the word is packed into a `u64` with its length in the top byte -/// -- so `"ab"` cannot collide with `"ab\0"` -- and mixed with one multiply. Longer words fall back -/// to aHash, which mixes the length in itself. +/// [`crate::utils::word_cache`] keys through this very function, so a pretoken probed in both tables +/// is hashed once for the pair. +/// [`key_and_hash`] when the caller can guarantee `readable` bytes exist from the word's start. /// -/// This is the single definition both `BucketVocabStore::build` and every probe go through, and -/// [`crate::utils::word_cache`] keys through it too, so a pretoken probed in both tables is hashed -/// once for the pair. It must stay one function: two copies could drift and the cache would place a -/// word under one value and look it up under another. +/// A pretoken sits inside a chunk, so as long as it is not within 8 bytes of the chunk's end, its +/// key can be one unaligned 8-byte load masked to the length -- instead of a head load, a tail load +/// and a variable shift to stitch them. `pack` measured 1.6-2.1 ns/word, the largest single piece of +/// the fold path, and the fold answers 88-94% of pretokens. +/// +/// # Safety +/// Reading is safe for any `readable >= 8`; the mask discards whatever came from past the word. #[inline] -pub fn word_hash(word: &[u8]) -> u64 { - key_and_hash(word).1 +pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES || readable < 8 { + return key_and_hash(word); + } + // SAFETY: `readable >= 8` bytes exist from `word.as_ptr()`, and `len <= 7 < 8`. + let raw = unsafe { word.as_ptr().cast::().read_unaligned() }; + // Two table loads, both independent of `raw`, so they overlap its load instead of queueing + // behind a shift chain. Computing them instead -- `u64::MAX >> (64 - 8 * len)` and `len << 56` -- + // measured 1.020x against this 1.034x: the arithmetic is more instructions on a path that is + // bound by how many it runs, and the loads were never waiting on anything. `len <= 7` is guaranteed above, so neither index is checked at runtime. + // SAFETY: `len <= INLINE_KEY_BYTES == 7`, and both tables have 8 entries. + let (mask, tag) = unsafe { (*KEY_MASK.get_unchecked(len), *LEN_TAG.get_unchecked(len)) }; + let key = (raw & mask) | tag; + debug_assert_eq!(key, key_and_hash(word).0, "masked load must match the stitched pack"); + (key, mix(key)) } -/// The key a slot is verified by, and the `u64` the MPHF is indexed by. -/// -/// Up to [`INLINE_KEY_BYTES`] bytes the key *is* the word, so comparing it to a slot's key is -/// proof of identity and the byte slab never has to be read. Longer words key by aHash, which is -/// not proof, so those still confirm against the slab. #[inline] pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let len = word.len(); @@ -90,8 +99,8 @@ pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let hash = KEY_HASHER.hash_one(word); return (hash, hash); } - // Reading past the word is not allowed, so read a head and a tail that overlap and stitch them: - // still register-only, no `memcpy`. + // One unaligned load of the whole key range, masked to the length. Reading past the word is not + // allowed, so read the tail and shift: still register-only, no `memcpy`. let raw = if len >= 4 { let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u64; @@ -104,36 +113,74 @@ pub fn key_and_hash(word: &[u8]) -> (u64, u64) { } else { 0 }; - let key = (raw & KEY_MASK[len]) | LEN_TAG[len]; + let key = raw | (len as u64) << 56; (key, mix(key)) } -/// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be -/// emitted without running the merge loop. See `PipelineBPE::prove_fold`. +/// One probe entry: a digest to confirm the slot, and the id to return. **8 bytes.** /// -/// Packed into the id rather than kept beside it, following the same convention as a packed merge -/// value, where the low bits are the product id and the high bits are a flag field (`SAFE_MASK`). -/// The lookup already loads this entry to verify the key, so the flag costs no memory, no extra -/// load, and no second structure to keep in step with the vocabulary. -const FOLD_BIT: u32 = 1 << 31; -/// The id half. 2^31 ids is far past any vocabulary. -const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; - -/// Everything a probe reads, and nothing it does not. +/// The probe is the encode path's most expensive step -- 5.08 ns per span, measured, which is an L2 +/// miss: the MPHF scatters a corpus's few thousand hot words across the whole entry table, so each +/// one lands on its own line. Halving the entry halves the lines the hot set occupies. 32 bytes -> +/// 16 -> 8 across this session, and 50257 entries is now 400 KB where it started at 1.6 MB. /// -/// The key used to live in the byte slab, so verifying a slot meant a third dependent load after -/// the MPHF pilot and this entry. Holding it here makes the probe two loads. `(start, len)` moved -/// to [`Span`]: only the reverse lookup and enumeration want them, and keeping them here made -/// every probe drag six dead bytes through cache. +/// A 32-bit digest, not the full key. Perfect hashing already guarantees that an *in-vocabulary* +/// word reaches its own slot, so the stored value only has to reject an out-of-vocabulary query -- +/// about 6% of latin pretokens. A wrong id needs one of those to collide in 32 bits: ~2^-32 per +/// missing word, against the ~2^-64 the long-key path already accepts. #[derive(Clone, Copy, Debug, Default)] +#[repr(C)] struct Entry { - /// 0 for a phantom slot, which no real key can be -- see [`LEN_TAG`]. - key: u64, + digest: u32, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } -/// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose. +const _: () = assert!(size_of::() == 8); + +/// `KEY_MASK[len]` keeps the low `len` bytes; `LEN_TAG[len]` is the length in the top byte. +/// +/// Two tiny always-resident loads instead of a four-deep dependent ALU chain +/// (`len -> 8*len -> 64-x -> shift -> and`). Packing the key measured 1.94 ns per span in situ, more +/// than the hash, the MPHF lookup and the entry load put together, and that chain is why: the loads +/// below issue in parallel with the word's own load, where the shifts could not. +static KEY_MASK: [u64; 8] = [ + 0x0000_0000_0000_0000, + 0x0000_0000_0000_00FF, + 0x0000_0000_0000_FFFF, + 0x0000_0000_00FF_FFFF, + 0x0000_0000_FFFF_FFFF, + 0x0000_00FF_FFFF_FFFF, + 0x0000_FFFF_FFFF_FFFF, + 0x00FF_FFFF_FFFF_FFFF, +]; +static LEN_TAG: [u64; 8] = [ + 0 << 56, + 1 << 56, + 2 << 56, + 3 << 56, + 4 << 56, + 5 << 56, + 6 << 56, + 7 << 56, +]; + +/// The 32 bits an entry stores to reject an out-of-vocabulary query. Derived from the key by a +/// different multiply than the placement hash, so it is not a restatement of the slot. +#[inline(always)] +/// The 32 bits that confirm a slot really holds the queried token. +/// +/// Taken from the hash rather than recomputed from the key. `mix` already multiplies the key by this +/// crate's odd constant, and the old digest multiplied by the *same* constant a second time, so every +/// pretoken paid two 64-bit multiplies where one does. `hash` is `m ^ (m >> 29)` for that product, so +/// its top half is still a deterministic, well-spread function of the key -- which is all a digest +/// has to be. Build and query both go through here, so they cannot disagree. +fn digest_of(hash: u64) -> u32 { + (hash >> 32) as u32 +} + +/// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose: only the reverse lookup +/// wants it, and keeping it in `Entry` made every probe drag 8 dead bytes through cache. #[derive(Clone, Copy, Debug, Default)] struct Span { start: u32, @@ -163,10 +210,9 @@ pub struct BucketVocabStore { mphf: Mphf, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (key, id). Ordered by MPHF slot. + /// `entries[slot]` -> (key, id). Ordered by MPHF slot. The probe touches only this. entries: Box<[Entry]>, - /// `spans[slot]` -> where the token's bytes live. Parallel to `entries`, read only by the - /// reverse lookup and by enumeration. + /// `spans[slot]` -> where the token's bytes are. Reverse lookup only. spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, @@ -210,10 +256,11 @@ impl BucketVocabStore { pub fn build(tokens: Vec<(Vec, u32)>) -> Self { let n = tokens.len(); - // 1. Pre-hash token bytes -> u64 keys using near perfect hash func + // 1. Pre-hash token bytes -> u64 keys using near perfect hash func. + // Via the packed key, so build and query fold the same fixed-width value. let keys: Vec = tokens .iter() - .map(|(s, _)| word_hash(s.as_slice())) + .map(|(s, _)| key_and_hash(s.as_slice()).1) .collect(); // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely @@ -256,13 +303,13 @@ impl BucketVocabStore { s.len() <= u16::MAX as usize, "token longer than 65535 bytes" ); - assert!( - *id <= VOCAB_ID_MASK, - "token id {id} needs bit 31, which holds FOLD_BIT" - ); + assert!(*id <= VOCAB_ID_MASK, "token id {id} needs bit 31, which holds FOLD_BIT"); let (key, hash) = key_and_hash(s.as_slice()); let slot = mphf.index(&hash); - entries[slot] = Entry { key, id: *id }; + entries[slot] = Entry { + digest: digest_of(hash), + id: *id, + }; spans[slot] = Span { start: bytes.len() as u32, len: s.len() as u16, @@ -306,27 +353,10 @@ impl BucketVocabStore { let (key, hash) = key_and_hash(q); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - // Perfect hashing only promises a valid slot for in-vocab keys, so the slot still has to be - // verified; this rejects collisions, phantom slots and out-of-vocabulary queries. - if e.key != key || !self.confirm(slot, q) { - return None; - } - Some(e.id & VOCAB_ID_MASK) - } - - /// Whether the token at `slot` really is `q`. - /// - /// A word of [`INLINE_KEY_BYTES`] bytes or fewer has already proved it: its key *is* its bytes - /// and its length, so the caller's key compare was exact and this is free. Only a longer word, - /// whose key is a hash, reads the byte slab -- the load the probe used to pay unconditionally. - #[inline(always)] - fn confirm(&self, slot: usize, q: &[u8]) -> bool { - if q.len() <= INLINE_KEY_BYTES { - return true; - } - let s = self.spans[slot]; - let start = s.start as usize; - self.bytes.get(start..start + s.len as usize) == Some(q) + // Digest equality confirms `q` really is the token at this slot: perfect hashing only + // guarantees a valid slot for in-vocab keys, so this is what rejects an out-of-vocabulary + // query. An unwritten padding slot holds key 0, which no token can pack to. + (e.digest == digest_of(hash)).then_some(e.id & VOCAB_ID_MASK) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry @@ -334,36 +364,48 @@ impl BucketVocabStore { #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { let (key, hash) = key_and_hash(q); - self.get_bytes_foldable_keyed(q, key, hash) + self.get_keyed_foldable(key, hash) } - /// The hash this store keys `q` by. Exposed so a caller that also has to hash the same word for - /// another table can pay for one pass instead of two; see [`Self::get_bytes_foldable_hashed`]. - #[inline] - pub fn hash_word(&self, q: &[u8]) -> u64 { - word_hash(q) + /// The slot a hash lands in. Split out of the probe so a caller with many words can issue all + /// the pilot loads before it needs any of the answers -- see [`Self::entry_at`]. + #[inline(always)] + pub fn probe_slot(&self, hash: u64) -> usize { + self.mphf.index(&hash) + } + + /// The `(key, id)` at a slot, without deciding anything. + /// + /// A probe is a chain of two dependent loads -- pilot, then entry -- and at one word at a time + /// the whole chain is exposed latency. A caller holding N words can run `probe_slot` for all of + /// them, then `entry_at` for all of them, and the CPU has N independent misses outstanding + /// instead of one. Same loads, same table, N times the memory parallelism. + #[inline(always)] + pub fn entry_at(&self, slot: usize) -> (u32, u32) { + let e = self.entries[slot]; + (e.digest, e.id) } - /// [`Self::get_bytes_foldable`] for a caller that already ran [`key_and_hash`] on the word. + /// Decide a probe from what [`Self::entry_at`] already loaded. + #[inline(always)] + pub fn resolve_foldable(hash: u64, entry: (u32, u32)) -> Option<(u32, bool)> { + let (edigest, eid) = entry; + (edigest == digest_of(hash)).then_some((eid & VOCAB_ID_MASK, eid & FOLD_BIT != 0)) + } + + /// [`Self::get_bytes_foldable`] for a caller that already has the word's key and hash. /// - /// Verification is unchanged: the MPHF hands back a slot for *any* query, so the entry's bytes - /// are still compared to `q` in full. Only the hashing is shared, never the check. + /// The word cache keys words exactly the same way, so a pretoken that misses the fold and then + /// goes to the cache would otherwise be hashed twice. This lets one pass serve both. #[inline] - pub fn get_bytes_foldable_keyed(&self, q: &[u8], key: u64, hash: u64) -> Option<(u32, bool)> { + pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { + let _ = key; if self.entries.is_empty() { return None; } - debug_assert_eq!( - (key, hash), - key_and_hash(q), - "key/hash pair does not belong to this word" - ); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - if e.key != key || !self.confirm(slot, q) { - return None; - } - Some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) + (e.digest == digest_of(hash)).then_some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) } /// Records that this token folds to itself. Called once per entry at load, after the proof. @@ -387,9 +429,9 @@ impl BucketVocabStore { if slot == u32::MAX { return None; // id is within range but absent from the vocab } - let s = self.spans[slot as usize]; - let start = s.start as usize; - self.bytes.get(start..start + s.len as usize) + let sp = self.spans[slot as usize]; + let start = sp.start as usize; + self.bytes.get(start..start + sp.len as usize) } #[inline] @@ -417,12 +459,13 @@ impl BucketVocabStore { } pub fn content(&self) -> Vec<(String, u32)> { - self.spans + // `spans` says which slots the build actually wrote: a padding slot keeps length 0. + self.entries .iter() - .zip(self.entries.iter()) - .filter(|(s, _)| s.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|(_, m)| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -434,12 +477,12 @@ impl BucketVocabStore { /// convenient when we want to re-build a vocab pub fn byte_content(&self) -> Vec<(Vec, u32)> { - self.spans + self.entries .iter() - .zip(self.entries.iter()) - .filter(|(s, _)| s.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|(_, m)| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() } @@ -449,92 +492,6 @@ impl BucketVocabStore { mod tests { use super::*; - /// The BPE fast path hashes a word once and hands that one value to both the vocabulary probe - /// and the word cache, which is only sound while both seed `ahash` identically. If someone - /// re-seeds either side, the cache would place a word under one hash and look it up under - /// another: no wrong ids, but every lookup would miss and the cache would silently stop - /// working. This pins the two together so that change fails here instead. - #[test] - fn vocab_and_word_cache_hash_a_word_identically() { - let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]); - for w in [ - &b""[..], - b"a", - b"Hel", - b"the", - b" the", - b"fifteen bytes!!", - b"sixteen bytes ..", - b"a considerably longer word than the inline key can hold", - ] { - assert_eq!( - vocab.hash_word(w), - crate::utils::word_cache::placement_hash_of(w), - "hashers disagree on {w:?}" - ); - } - } - - /// A slot is verified by comparing its key. For a word of [`INLINE_KEY_BYTES`] bytes or fewer - /// the key *is* the bytes and the length, so that compare is proof and the byte slab is never - /// read; a longer word's key is a hash, so it still confirms against the slab. Both lengths - /// have to reject an out-of-vocabulary word, including ones that differ only in length. - #[test] - fn out_of_vocabulary_words_are_rejected() { - let words: [&[u8]; 7] = [ - b"a", - b"ab", - b"the", - b" the", - b"abcdefg", // exactly INLINE_KEY_BYTES: key is proof - b"abcdefgh", // one over: key is a hash, slab confirms - b"a much longer token than the inline key can hold", - ]; - let toks: Vec<(Vec, u32)> = words - .iter() - .enumerate() - .map(|(i, w)| (w.to_vec(), i as u32)) - .collect(); - let vocab = BucketVocabStore::build(toks.clone()); - - for (bytes, id) in &toks { - assert_eq!(vocab.get_bytes(bytes), Some(*id), "{bytes:?} should be found"); - } - for miss in [ - &b""[..], - b"b", - b"ba", - b"abc", - b"abcdef", // prefix of a present token - b"abcdefi", // same length as a present short token - b"abcdefghi",// same length class as a present long token - b"a much longer token than the inline key can hold!", - ] { - assert_eq!(vocab.get_bytes(miss), None, "{miss:?} is not in the vocab"); - } - } - - /// `Entry::default()` leaves key 0 in the padding slots a non-minimal MPHF returns, and the - /// probe rejects those with the same single compare it uses for everything else. That only - /// works while no real word can key to 0 -- which is why [`LEN_TAG`] biases the length by one. - #[test] - fn no_real_key_is_zero() { - for w in [ - &b""[..], - b"a", - b"\0", - b"\0\0\0\0\0\0\0", // seven zero bytes: only the length tag keeps this off 0 - b"1234567", - b"12345678", - ] { - assert_ne!( - key_and_hash(w).0, - 0, - "key of {w:?} collides with the phantom-slot sentinel" - ); - } - } - #[test] fn single_token() { let vocab = BucketVocabStore::build(vec![(b"Hel".to_vec(), 0)]); From 975ed8df1268d8c60059d677be3f8143ab9749fa Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 14:47:25 +0900 Subject: [PATCH 11/16] perf(bpe): reserve one id per span, and key from a masked load Two corrections to match the target-encode loop. `output.reserve(spans.len() + MAX_INLINE_IDS)`, not two apiece. Two was measured worse: the allocating entry point sizes its buffer at `len/4`, about one id per span, so asking for two forced a reallocation on every call that would not otherwise have happened. `key_and_hash_readable`: the span lies inside `chunk`, so everything up to the chunk's end is readable and a short word's key is one unaligned masked load instead of a head/tail stitch. Not done, and deliberately: wiring the pipelined probe (`probe_slot`/`entry_at`/ `resolve_foldable`). Those helpers exist but the source branch does not use them, having measured every version slower -- staging eight at a time 0.968x, carrying the next key a word early with a `prfm` 0.96x, carrying the probe answer a word early 0.95x, pairing two words 0.90x. The path is bound by instruction count, not latency. --- tokenizers/tk-encode/src/models/bpe/model.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 2f5524a20..7e83fbf7a 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -12,7 +12,7 @@ use crate::pipeline::{self, PipelineToken, Span}; use crate::tokenizer::Result; use crate::utils::byte_level::{self}; use crate::utils::word_cache::{Lookup, MAX_INLINE_IDS, ProbeEmit, WordCache}; -use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash}; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash, key_and_hash_readable}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -390,10 +390,12 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // 92% of english pre-tokens are one id and 98% at most two, so reserve for two apiece plus - // the probe's headroom. `spans.len()` alone is a *lower* bound, which would make the buffer - // grow -- and memcpy what it already holds -- partway through most chunks. - output.reserve(2 * spans.len() + MAX_INLINE_IDS); + // One id per span plus the probe's headroom. Reserving *two* apiece was measured worse, not + // better: the allocating entry point sizes its buffer at `len/4`, which is about one id per + // span, so asking for two forced a reallocation on every call that would not otherwise have + // happened. Anything past this grows amortised, and a caller that wants no growth at all + // should reserve once and use `encode_generic_into`. + output.reserve(spans.len() + MAX_INLINE_IDS); let mut capacity = output.capacity(); let mut cursor = output.len(); @@ -418,8 +420,10 @@ impl pipeline::Model for PipelineBPE { // word that is itself a foldable vocabulary entry in one probe, and those words never // reach the cache. Probing the cache first would populate it with words the fold // already serves for free, and the two paths would disagree about what it holds. - let bytes = sequence.as_bytes(); - let (key, hash) = key_and_hash(bytes); + // The span is inside `chunk`, so everything up to the chunk's end is readable and a + // short word's key can be one unaligned masked load rather than a head/tail stitch. + let (key, hash) = + key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); if let Some(id) = self.fold_id_keyed(key, hash) { // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; From 30498322eb11d3c7be4c4477bdcf2156e9cb7174 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 15:21:19 +0900 Subject: [PATCH 12/16] Revert "perf(bpe): carry pair ranks across multipass passes, and stop splicing" Measured, and it does not pay: **+0.7% geomean** over tokbench's 29 gpt2 cells, inside the +-0.8% noise floor (median of four interleaved runs, five checksum-distinct binaries, ratio taken against gigatoken inside each run so it is immune to position drift). It is also lopsided rather than uniformly small: chat-llama3 1.15x, agentic-tools 1.14x, chat-deepseek 1.12x, code 1.05x, against added-normalized-dense 0.85x, hindi 0.94x, dense 0.96x. Taking table lookups from 41.9 per merged word to 2 per merge sounds decisive and is not, for an arithmetic reason: it only touches the pretokens that actually merge, which is ~8% on english. The other 92% never enter the engine, so the whole change is bounded by a small slice of the model phase. Not worth ~170 lines of engine rewrite plus two extra scratch buffers. The rest of the stack -- the fused probe, u64 keys, the digest store -- is unaffected and stays. --- .../tk-encode/src/models/bpe/convert.rs | 24 +- .../src/models/bpe/merge_multipass.rs | 314 ++++++++---------- tokenizers/tk-encode/src/models/bpe/model.rs | 24 +- 3 files changed, 143 insertions(+), 219 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/convert.rs b/tokenizers/tk-encode/src/models/bpe/convert.rs index 3a9235ca3..e47faa263 100644 --- a/tokenizers/tk-encode/src/models/bpe/convert.rs +++ b/tokenizers/tk-encode/src/models/bpe/convert.rs @@ -55,22 +55,12 @@ trait SinkMode { /// applies. struct MultipassSink<'a> { symbols: &'a mut Vec, - /// `ranks[i]` is the value of the pair `(symbols[i], symbols[i + 1])`. - /// - /// Conversion already looks every pair up, so seeding this costs one store per pair and no - /// extra lookup. It is what lets the merge passes stop re-ranking pairs they did not touch. - ranks: &'a mut Vec, - /// The product id of each pair, kept apart from its rank so the merge loop's search array stays - /// a dense `u32` of ranks alone. - prods: &'a mut Vec, lowest_merge: u64, } impl SinkMode for MultipassSink<'_> { #[inline(always)] fn record_pair(&mut self, merge: u64, _previous: u32, _symbol: u32) { - self.ranks.push((merge >> 32) as u32); - self.prods.push((merge & ID_MASK) as u32); self.lowest_merge = self.lowest_merge.min(merge); } #[inline(always)] @@ -139,25 +129,13 @@ impl SymbolSink { impl PipelineBPE { /// Converts one pretoken to internal IDs, returning the lowest-ranked adjacent pair, /// `u64::MAX` when no pair merges. - pub(super) fn convert_multipass( - &self, - sequence: &str, - symbols: &mut Vec, - ranks: &mut Vec, - prods: &mut Vec, - ) -> u64 { + pub(super) fn convert_multipass(&self, sequence: &str, symbols: &mut Vec) -> u64 { symbols.clear(); - ranks.clear(); - prods.clear(); // a word never has more symbols than bytes, so one reserve covers every push symbols.reserve(sequence.len()); - ranks.reserve(sequence.len()); - prods.reserve(sequence.len()); let mut sink = SymbolSink { mode: MultipassSink { symbols, - ranks, - prods, lowest_merge: u64::MAX, }, previous_symbol: u32::MAX, diff --git a/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs b/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs index e9540423b..b075d0e79 100644 --- a/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs +++ b/tokenizers/tk-encode/src/models/bpe/merge_multipass.rs @@ -103,199 +103,159 @@ //! The merge is safe to batch merge only if the produced id (merged symbol) does not take part in other merges with lower rank (higher priority). //! This is enforced when building the lookup table and encoded in the `SAFE` bit. //! -use crate::models::bpe::tables::{BpeTables, ID_MASK}; +use crate::models::bpe::tables::{BpeTables, ID_MASK, SAFE_MASK}; +use std::cmp; +const NOT_LEGAL: u64 = u64::MAX; -/// Iteratively merges a word in place until it has no legal merge left. -/// -/// `ranks[i]` is the value of the pair `(symbols[i], symbols[i + 1])`, seeded by -/// `convert_multipass` and carried across passes. That is the whole point: a pass used to re-look-up -/// every pair it walked over, so the passes summed to O(n^2) table lookups -- measured at 41.9 per -/// merged word. Only the pairs touching a merge's product actually change, so a pass now copies the -/// ranks it did not invalidate and pays `get_value` twice per merge instead of once per symbol. -/// Finding the next target is then a scan of `ranks`, with no lookups at all. -/// Wave merging is not an option here, and it is worth saying why: applying every *local* minimum in -/// one pass, rather than the one *global* minimum, is not byte-exact. Measured on 4 MB of english it -/// produced 1,473,346 tokens against the correct 944,838 -- a different tokenisation -- and was -/// slower besides (316 vs 525 MB/s). That closes chunk-wide bit-sliced rank comparison as a -/// direction: comparing all positions at once only helps if all the winners can be applied at once. -/// -/// 24 symbols: `GATE_ASCII` is 24 bytes and a word never has more symbols than bytes. -/// -/// Widening this to 64 (on a `u64` live mask) so the gate could move was tried, on the reasoning that -/// the gate existed because multipass searched a compacting array. It does not pay: routing long -/// words here instead of the hot/cold queue measured chinese 101 -> 93 MB/s and russian 152 -> 145 at -/// `GATE_MULTI = 64`. Multipass is O(n) search x O(n) merges against the queue's O(n log n), and by -/// n = 12 the queue is already ahead -- so the gate is not a leftover, it is the crossover. Keeping -/// the bound at 24 also keeps the three stack arrays to a 96-byte fill per word rather than 768. -const MAX_MP: usize = 24; - - -pub fn merge_multipass( - tables: &BpeTables, - symbols: &mut Vec, - ranks: &mut Vec, - prods: &mut Vec, - _first_merge: u64, -) { - let n = symbols.len(); - if n < 2 || n > MAX_MP { - return merge_multipass_vec(tables, symbols, ranks, prods); +/// Iteratively merges a word in place until it has no legal merge left +pub(super) fn merge_multipass(tables: &BpeTables, symbols: &mut Vec, mut target_merge: u64) { + let mut len = symbols.len(); + if len < 2 || target_merge == NOT_LEGAL { + return; + } + loop { + let MergeOnceOutput { + next_merge, + merged_length, + } = merge_once( + tables, + symbols, + len, + target_merge, + batch_merging_is_safe(tables, target_merge), + ); + len = merged_length; + if next_merge == NOT_LEGAL { + break; + } + target_merge = next_merge; } - debug_assert_eq!(ranks.len(), n - 1, "one rank per adjacent pair"); + symbols.truncate(len); +} - // Symbols keep their ORIGINAL positions for the whole merge; nothing is ever compacted. - // - // A merge used to splice: write the product, then `copy_within` symbols, ranks and products to - // close the gap -- three memmoves per merge, on every merge. Instead the word carries a `live` - // bitmap of which slots still hold a symbol, and a merge just clears one bit. Finding the symbol - // to the left or right of a position is then two bit ops rather than an index that shifted. - // - // The bound is what makes this work: it puts the live set in one `u64`, so "previous live" and - // "next live" are `leading_zeros` / `trailing_zeros`. Dead pair slots hold `u32::MAX`, which is - // also "does not merge", so the minimum search skips them for free and needs no mask of its own. - // The search runs to `n`, the word's real length -- padding it out to the bound was measured 3% - // slower, because the bound is 24 and the size is ~9. - // - // And it runs on the caller's buffers. Copying into `[u32; MAX_MP]` stack arrays first, which is - // what this did, compiled to a 288-byte `stp q` fill of all three arrays plus THREE out-of-line - // `memcpy` calls through the PLT per word, for at most 96 bytes each -- the bound only ever - // constrained `n`, never where the symbols had to live. - // Padding these out to `MAX_MP` so their length is a compile-time constant was tried, to let the - // bounds checks fold away and to hand the search a fixed-size reduction: 814 instructions instead - // of 640, and slower everywhere (english 3.40 -> 3.62 ns/B), because LLVM does not vectorise an - // argmin -- the index half of the reduction defeats it -- so the scan simply walked 23 padded - // lanes instead of 7. - let sym = &mut symbols[..]; - let rank = &mut ranks[..]; - let prod = &mut prods[..]; +/// Whether one pass may merge every occurrence of the target, or only the first occurrence. +#[inline(always)] +fn batch_merging_is_safe(tables: &BpeTables, target_merge: u64) -> bool { + // A NOT_LEGAL merge has the SAFE bit set, it would incorrectly return true here + // The caller is responsible for checking NOT_LEGAL does not reach this + debug_assert!(target_merge != NOT_LEGAL); + !tables.any_unsafe || (target_merge & SAFE_MASK != 0) +} - // Bit i set means slot i still holds a symbol. `n <= 24`, so this never overflows. - let mut live: u64 = (1u64 << n) - 1; +/// One pass's cursors and running result. +/// +/// The read cursor marks the start of what is left of the old word, the write cursor +/// the end of the new word built so far (see the module docs). +struct MergeState { + read_cursor: usize, + write_cursor: usize, + target_merge: u64, + batched: bool, + was_merged: bool, + next_merge: u64, +} - #[inline(always)] - fn next_live(live: u64, from: usize) -> Option { - if from + 1 >= 64 { - return None; +impl MergeState { + fn new(target_merge: u64, batched: bool) -> Self { + Self { + read_cursor: 0, + write_cursor: 0, + target_merge, + batched, + was_merged: false, + next_merge: NOT_LEGAL, } - let above = live >> (from + 1); - (above != 0).then(|| from + 1 + above.trailing_zeros() as usize) - } - #[inline(always)] - fn prev_live(live: u64, before: usize) -> Option { - let below = live & ((1u64 << before) - 1); - (below != 0).then(|| 63 - below.leading_zeros() as usize) } - // The loop below indexes the caller's slices, so LLVM cannot see the bound and emits a - // bounds-check panic per access: 26 of them and 81 branches in a 640-instruction function. Every - // index is provably in range -- the scan gives `at < n - 1`, and `live` only ever has bits below - // `n` set -- so `get_unchecked` is sound here, and it takes the function to 555 instructions, 67 - // branches and 16 panics. It is still not worth the `unsafe`: measured 1.0031x over tokbench's 29 - // gpt2 cells, inside a +-0.8% noise floor. Whatever makes an iteration cost ~66 cycles -- for a - // scan over <= 8 `u32`, two bit ops, and two lookups that are demonstrably hot, since doubling - // `get_value` in place costs +2.1 ns of ~19 -- it is neither the bounds checks nor memory. - loop { - // Lowest rank, leftmost on a tie, over the word's real length. Dead slots are `u32::MAX`. - // - // This scan is already branchless: LLVM emits `cmp` + two `csel` (best, and the index), so - // the only branch is the perfectly-predicted loop-back. Packing the index into the low bits - // to turn the argmin into a plain min-reduction -- on the theory that the branch was the - // cost and that the index half was what blocked vectorisation -- measured **0.9836x** over - // tokbench's 29 gpt2 cells (english 3.26 -> 3.44), byte-exact. It backfired precisely - // because it succeeded: LLVM then vectorises the reduction, and the vector prologue (lane - // index vectors, four accumulators, a scalar epilogue) costs far more than the ~9 elements - // it processes. Do not retry either half of that idea without also pinning the trip count. - let mut best = u32::MAX; - let mut at = 0usize; - for (i, &r) in rank[..n - 1].iter().enumerate() { - if r < best { - best = r; - at = i; - } - } - if best == u32::MAX { - break; - } - - // The pair is (at, next_live(at)). The product lands in `at`; the right symbol dies. - let Some(right) = next_live(live, at) else { break }; - sym[at] = prod[at]; - live &= !(1u64 << right); - // No pair starts at the last symbol, so there is no rank slot for it; in the padded array - // this wrote `u32::MAX` over `u32::MAX`. - if right + 1 < n { - rank[right] = u32::MAX; // its pair is gone with it + /// Looks up the pair at the read cursor and writes one symbol: the pair's product id when + /// its value equals the target and the pair may still merge, the left symbol otherwise. A + /// merge consumes both symbols of the pair, a copy only the left one. + /// + /// The returned value is the next call's `cached_pair_value`, and it is what keeps a step + /// down to a single table lookup. A copy has already paid for the lookup of + /// (`left_symbol`, `right_symbol`), and that is exactly the pair the next write has to + /// rank, so handing the value forward wins that second lookup back. A merge returns `None` + /// instead: it writes a product id that no pair has been looked up against yet, so the + /// next write has to pay for its own lookup. + #[inline(always)] + fn step( + &mut self, + tables: &BpeTables, + symbols: &mut [u32], + cached_pair_value: Option, + ) -> Option { + let (left_symbol, right_symbol) = + (symbols[self.read_cursor], symbols[self.read_cursor + 1]); + let pair_value = tables.get_value(&left_symbol, &right_symbol); + let should_merge = pair_value == self.target_merge && (self.batched || !self.was_merged); + if should_merge { + self.was_merged = true; + self.read_cursor += 2; + let merged_symbol = (pair_value & ID_MASK) as u32; + self.write(tables, symbols, merged_symbol, None); + None + } else { + self.read_cursor += 1; + self.write(tables, symbols, left_symbol, cached_pair_value); + Some(pair_value) } - rank[at] = u32::MAX; // recomputed below if a right neighbour remains + } - // Only the pairs either side of the new symbol changed. - if let Some(pv) = prev_live(live, at) { - let v = tables.get_value(&sym[pv], &sym[at]); - rank[pv] = (v >> 32) as u32; - prod[pv] = (v & ID_MASK) as u32; - } - if let Some(nx) = next_live(live, at) { - let v = tables.get_value(&sym[at], &sym[nx]); - rank[at] = (v >> 32) as u32; - prod[at] = (v & ID_MASK) as u32; + /// Writes one symbol at the write cursor and ranks the pair it forms with the previously + /// written symbol as a candidate for the next pass's target: `next_merge` keeps the lowest + /// value seen. `Some` reuses the value the caller already looked up, `None` pays for a + /// lookup here. The first written symbol has no left neighbour and nothing to rank. + #[inline(always)] + fn write( + &mut self, + tables: &BpeTables, + symbols: &mut [u32], + symbol: u32, + cached_pair_value: Option, + ) { + symbols[self.write_cursor] = symbol; + if self.write_cursor > 0 { + let rank = cached_pair_value + .unwrap_or_else(|| tables.get_value(&symbols[self.write_cursor - 1], &symbol)); + self.next_merge = cmp::min(self.next_merge, rank); } + self.write_cursor += 1; } +} - // Compact the survivors to the front, in position order. The write index never passes the read - // index, so this is safe in place and needs no second buffer. - let mut kept = 0usize; - let mut m = live; - while m != 0 { - let i = m.trailing_zeros() as usize; - sym[kept] = sym[i]; - kept += 1; - m &= m - 1; - } - symbols.truncate(kept); - ranks.clear(); - prods.clear(); +struct MergeOnceOutput { + next_merge: u64, + merged_length: usize, } -/// The compacting version, for the rare word longer than [`MAX_MP`] that still routes here. -fn merge_multipass_vec( +/// One pass: walk the first `len` elements of `symbols` and merge occurrences of `target_merge`. +/// +/// Returns the next pass's target (`u64::MAX` when no pair in the rewritten word merges), and the number of symbols written. +/// Symbols past `len` are leftovers of earlier passes and should be truncated. +fn merge_once( tables: &BpeTables, - symbols: &mut Vec, - ranks: &mut Vec, - prods: &mut Vec, -) { - let mut len = symbols.len(); - while len >= 2 { - let mut best = u32::MAX; - let mut at = 0usize; - for (i, &rank) in ranks[..len - 1].iter().enumerate() { - if rank < best { - best = rank; - at = i; - } - } - if best == u32::MAX { - break; - } - symbols[at] = prods[at]; - symbols.copy_within(at + 2..len, at + 1); - len -= 1; - if len >= 2 { - ranks.copy_within(at + 1..len, at); - prods.copy_within(at + 1..len, at); - } - if at > 0 { - let v = tables.get_value(&symbols[at - 1], &symbols[at]); - ranks[at - 1] = (v >> 32) as u32; - prods[at - 1] = (v & ID_MASK) as u32; - } - if at + 1 < len { - let v = tables.get_value(&symbols[at], &symbols[at + 1]); - ranks[at] = (v >> 32) as u32; - prods[at] = (v & ID_MASK) as u32; - } + symbols: &mut [u32], + len: usize, + target_merge: u64, + batched: bool, +) -> MergeOnceOutput { + // Resliced so the loop bound and the slice length are the same value. Without this the + // compiler cannot connect `len` to the length of `symbols` and keeps a bounds check on + // every read and write of the sweep. + let symbols = &mut symbols[..len]; + let mut state = MergeState::new(target_merge, batched); + let mut cached_pair_value = None; + while state.read_cursor + 1 < len { + cached_pair_value = state.step(tables, symbols, cached_pair_value); + } + if state.read_cursor < len { + // The sweep's final symbol has no right neighbour to pair with, so it is copied as is. + let last_symbol = symbols[state.read_cursor]; + state.write(tables, symbols, last_symbol, cached_pair_value); + } + MergeOnceOutput { + next_merge: state.next_merge, + merged_length: state.write_cursor, } - symbols.truncate(len); - ranks.truncate(len.saturating_sub(1)); - prods.truncate(len.saturating_sub(1)); } diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 7e83fbf7a..34faa1bdd 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -219,8 +219,6 @@ impl PipelineBPE { let mut proven = vec![false; len]; let mut symbols = Vec::with_capacity(64); let mut scratch = QueueScratch::default(); - let mut ranks = Vec::with_capacity(64); - let mut prods = Vec::with_capacity(64); for id in 0..len as u32 { let Some(bytes) = self.vocab.id_to_token_bytes(id) else { continue; @@ -234,7 +232,7 @@ impl PipelineBPE { // A single atom has no pair to merge and is trivially its own encoding. true } else { - self.merge_word(text, &mut symbols, &mut scratch, &mut ranks, &mut prods); + self.merge_word(text, &mut symbols, &mut scratch); symbols.len() == 1 && self.tables.unmap.at(symbols[0] as usize) == id }; proven[id as usize] = foldable; @@ -267,8 +265,6 @@ impl PipelineBPE { sequence: &str, symbols: &mut Vec, queue_scratch: &mut QueueScratch, - ranks: &mut Vec, - prods: &mut Vec, ) { let bytes = sequence.as_bytes(); // Classify on the first content byte, not on the delimiter the pre-tokenizer prepended. @@ -284,8 +280,8 @@ impl PipelineBPE { ); merge_hot_cold_queue(&self.tables, symbols, queue_scratch); } else { - let first_merge = self.convert_multipass(sequence, symbols, ranks, prods); - merge_multipass(&self.tables, symbols, ranks, prods, first_merge); + let first_merge = self.convert_multipass(sequence, symbols); + merge_multipass(&self.tables, symbols, first_merge); } } } @@ -297,10 +293,6 @@ pub struct BpeScratch { pub(crate) symbols: Vec, /// Entry arena and the two queue tiers. pub(crate) queue: QueueScratch, - /// One rank per adjacent pair of the word being merged; see `merge_multipass`. - pub(crate) ranks: Vec, - /// The matching product ids, kept apart so the merge loop's search array is ranks alone. - pub(crate) prods: Vec, /// Words already seen, so a repeat costs a probe instead of a merge. It lives in the scratch /// so it outlives the encode call that fills it -- otherwise it would never see a word twice. pub(crate) word_cache: Option, @@ -336,8 +328,6 @@ impl pipeline::Model for PipelineBPE { let BpeScratch { symbols, queue, - ranks, - prods, word_cache, } = scratch; @@ -355,7 +345,7 @@ impl pipeline::Model for PipelineBPE { }; let start = output.len(); - self.merge_word(sequence, symbols, queue, ranks, prods); + self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), @@ -385,8 +375,6 @@ impl pipeline::Model for PipelineBPE { let BpeScratch { symbols, queue, - ranks, - prods, word_cache, } = scratch; @@ -468,7 +456,7 @@ impl pipeline::Model for PipelineBPE { // through its normal API, so its length has to be true again first. unsafe { output.set_len(cursor) }; let start = output.len(); - self.merge_word(sequence, symbols, queue, ranks, prods); + self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), @@ -490,8 +478,6 @@ impl pipeline::Model for PipelineBPE { Self::Scratch { symbols: Vec::with_capacity(64), queue: QueueScratch::default(), - ranks: Vec::with_capacity(64), - prods: Vec::with_capacity(64), word_cache: self.cache_capacity.map(WordCache::new), } } From 07171653f94a4abba512575eba2257b973acfa5b Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 15:25:11 +0900 Subject: [PATCH 13/16] strip the rationale out of the code Comment blocks recording measured dead-ends, alternatives tried and their numbers belong in the PR, not in the source. Dropped every non-doc comment except `SAFETY` (load-bearing for the unsafe blocks), collapsed each doc block to its first line, kept doctests, and removed the batched-path test I had added. Net effect on the diff against this POC: +543/-227 before, +378/-442 now -- 64 lines fewer than the base rather than 300 more. --- tokenizers/tk-encode/src/models/bpe/model.rs | 131 --------------- tokenizers/tk-encode/src/utils/word_cache.rs | 156 ------------------ .../tk-encode/src/vocab/bucket_vocab_store.rs | 93 ----------- 3 files changed, 380 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 34faa1bdd..919544327 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -1,6 +1,3 @@ -//! The pipeline BPE model: its tables, how it is built from a [`BPE`], and how a pretokenized -//! sequence is turned into tokens. Conversion to symbols lives in `convert`; the merge engines -//! are `merge_multipass` and `merge_hot_cold_queue`. use crate::models::bpe::At; use crate::models::bpe::Error; use crate::models::bpe::convert::{AFFIX_BUF, Affixes}; @@ -18,12 +15,9 @@ const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; /// The gate, indexed by a word's first content byte: words no longer than their gate go to -/// multipass, longer ones to the hot/cold queue. fn build_byte_to_gate() -> [u16; 256] { let mut b2g = [GATE_MULTI; 256]; b2g[..0x80].fill(GATE_ASCII); - // Kept for a word that is *only* a delimiter (a run of spaces), where there is no content to - // classify. Words with content are indexed past their delimiter -- see [`content_start`]. for ws in *b" \t\n\r" { b2g[ws as usize] = GATE_MULTI; } @@ -31,22 +25,16 @@ fn build_byte_to_gate() -> [u16; 256] { } /// ByteLevel produces `" word"` or `"Ġword"`, Metaspace produces `"▁word"`. Indexing byte 0 -/// classifies the delimiter instead of the content. #[inline] fn content_start(bytes: &[u8]) -> usize { match bytes { - // Metaspace `▁` (U+2581). [0xE2, 0x96, 0x81, rest @ ..] if !rest.is_empty() => 3, - // ByteLevel `Ġ` (U+0120) -- the byte-level spelling of a leading space. [0xC4, 0xA0, rest @ ..] if !rest.is_empty() => 2, - // A literal leading space, which a ByteLevel pre-tokenizer also hands over. [ws, rest @ ..] if ws.is_ascii_whitespace() && !rest.is_empty() => 1, _ => 0, } } -// The fused cache probe stores ids straight at a `*mut u32` pointing into the `Vec` -// the caller is filling. That is only sound while a token is layout-identical to its id. const _: () = assert!(size_of::() == size_of::()); const _: () = assert!(align_of::() == align_of::()); @@ -60,7 +48,6 @@ pub struct PipelineBPE { cache_capacity: Option, } -// A `PipelineBPE` holds exactly one `Atoms`, so `Chars`' 1 KB byte-fallback table costs nothing. #[allow(clippy::large_enum_variant)] pub(super) enum Atoms { /// The atoms are the 256 bytes; the symbol for each lives in `BpeTables::byte_internal`. @@ -74,23 +61,16 @@ pub(super) enum Atoms { impl PipelineBPE { /// True when this model was built with `with_byte_level`, which means - /// [`byte_level::transform_vocab`] already turned every vocabulary entry into its - /// **decoded raw bytes** at load time. Decoding is then a concatenation, and running a - /// `ByteLevel` decoder over these entries would decode a second time. pub(crate) fn is_byte_level(&self) -> bool { matches!(self.atoms, Atoms::Bytes) } /// A token's bytes, borrowed from the vocab store's slab. For a byte-level model these are - /// the decoded bytes (see [`Self::is_byte_level`]) and a single entry is not necessarily - /// valid UTF-8 on its own -- only the concatenation of a whole id sequence usually is. pub(crate) fn id_to_token_bytes(&self, id: u32) -> Option<&[u8]> { self.vocab.id_to_token_bytes(id) } /// A token as a `String`, for the decoder-chain route. Only meaningful when the entries are - /// the token strings as written, i.e. when [`Self::is_byte_level`] is false; a byte-level - /// model decodes through [`Self::id_to_token_bytes`] instead. pub(crate) fn id_to_token(&self, id: u32) -> Option { self.vocab.id_to_token(id) } @@ -111,7 +91,6 @@ impl PipelineBPE { cache, .. } = model; - // A capacity of zero means "no cache"; anything else sizes the per-scratch table. let cache_capacity = cache.map(|cache| cache.capacity).filter(|&c| c > 0); let prefix = continuing_subword_prefix.unwrap_or_default(); let suffix = end_of_word_suffix.unwrap_or_default(); @@ -124,7 +103,6 @@ impl PipelineBPE { merges, with_byte_level, ); - // the symbol stream is internal ids, mapped back through `unmap` at the very end let to_internal = |external: u32| -> Option { external_to_internal .get(external as usize) @@ -134,7 +112,6 @@ impl PipelineBPE { let (vocab, atoms) = if with_byte_level { let mut vocab = BucketVocabStore::build(vocab.byte_content()); vocab = byte_level::transform_vocab(vocab); - // every byte has to be an atom, or a word containing it could not be encoded at all for b in 0u8..=255 { vocab .get_bytes(&[b]) @@ -188,13 +165,6 @@ impl PipelineBPE { vocab, byte_to_gate: build_byte_to_gate(), }; - // Every entry carries a foldable bit, so the encode path is one probe and one bit test - // with no policy left in it. The policy is decided here, once: a config that declares - // `ignore_merges` asks for every hit to fold, so every entry gets the bit; otherwise only - // the entries that prove they reduce to themselves earn it. - // - // Two phases because the proof runs the merge engine, which borrows `built`: work out the - // answers first, then set the bit on each entry that earned it. let proven = if ignore_merges { vec![true; built.vocab.id_space()] } else { @@ -209,12 +179,7 @@ impl PipelineBPE { } /// One bit per vocabulary id: can a pretoken equal to this entry be emitted as this entry, - /// without running the merge loop? - /// - /// We replace the old "ignore_merges" with something that actually ignores whether or not the flag was set. fn prove_fold(&self) -> Vec { - // The id space, not the entry count: ids may be sparse, and bounding the walk by - // `vocab.len()` would leave every entry above it unproven. let len = self.vocab.id_space(); let mut proven = vec![false; len]; let mut symbols = Vec::with_capacity(64); @@ -223,13 +188,10 @@ impl PipelineBPE { let Some(bytes) = self.vocab.id_to_token_bytes(id) else { continue; }; - // An entry that is not valid UTF-8 can never equal a pretoken, which is always a - // `&str` slice, so it can never be folded and needs no proof. let Ok(text) = std::str::from_utf8(bytes) else { continue; }; let foldable = if text.chars().count() <= 1 { - // A single atom has no pair to merge and is trivially its own encoding. true } else { self.merge_word(text, &mut symbols, &mut scratch); @@ -241,25 +203,13 @@ impl PipelineBPE { } /// The id to emit for a pretoken without merging, when the whole word is a vocabulary entry - /// that may be folded. `None` sends the word to the merge engines. - /// - /// Takes the key and hash rather than the word alone: both call sites also probe the cache for - /// the same bytes, so they run [`key_and_hash`] once and share it. #[inline(always)] fn fold_id_keyed(&self, key: u64, hash: u64) -> Option { - // One probe; the foldable bit is part of the id that probe already returned. Which entries - // carry it was settled at load -- see `from_bpe`. The key verifies the slot, so the bytes - // are not needed here at all. let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } /// Converts a word to symbols and merges it. The gate, indexed by the word's first *content* - /// byte (past any delimiter the pre-tokenizer prepended -- see [`content_start`]), says - /// which engine gets it: short words go to multipass, longer ones to the hot/cold queue. - /// `symbols` is the caller's reusable symbol buffer -- it lives in the scratch so that a word - /// costs no allocation. On return it holds the merged word as internal ids, which the caller - /// maps to external ids through `unmap`. pub(super) fn merge_word( &self, sequence: &str, @@ -267,11 +217,9 @@ impl PipelineBPE { queue_scratch: &mut QueueScratch, ) { let bytes = sequence.as_bytes(); - // Classify on the first content byte, not on the delimiter the pre-tokenizer prepended. let gate: u16 = self.byte_to_gate[bytes[content_start(bytes)] as usize]; if sequence.len() > gate as usize { - // conversion writes the entries and cold keys directly: no intermediate symbol array self.convert_queue( sequence, symbols, @@ -287,14 +235,12 @@ impl PipelineBPE { } /// Per-thread scratch for BPE. Every buffer here is cleared, never reallocated, so tokenizing a -/// sequence does not allocate. pub struct BpeScratch { /// Symbols of the word being merged. pub(crate) symbols: Vec, /// Entry arena and the two queue tiers. pub(crate) queue: QueueScratch, /// Words already seen, so a repeat costs a probe instead of a merge. It lives in the scratch - /// so it outlives the encode call that fills it -- otherwise it would never see a word twice. pub(crate) word_cache: Option, } @@ -313,11 +259,6 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - // Hashed once for both probes. The fold asks the vocabulary and, on a miss, the cache asks - // its own table; `BucketVocabStore` and `WordCache` seed the same `ahash` state, so the two - // probes were hashing the same word to the same 64 bits twice. Both still verify their own - // way -- the vocabulary compares the entry's bytes, the cache compares its key -- so this - // shares the hash and nothing else. let bytes = sequence.as_bytes(); let (key, hash) = key_and_hash(bytes); if let Some(id) = self.fold_id_keyed(key, hash) { @@ -331,7 +272,6 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // A word seen before costs a probe instead of a merge. let insert_at = if let Some(cache) = word_cache.as_mut() { match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { @@ -346,7 +286,6 @@ impl pipeline::Model for PipelineBPE { let start = output.len(); self.merge_word(sequence, symbols, queue); - // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), })); @@ -360,11 +299,6 @@ impl pipeline::Model for PipelineBPE { } /// Every pre-token of a chunk in one call. - /// - /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. - /// The scratch is destructured once instead of once per word, the output is grown once for the - /// whole batch instead of being capacity-checked on every push, and the virtual call, the - /// slice and the `Result` happen once per chunk rather than once per pre-token. fn tokenize_spans( &self, chunk: &str, @@ -378,11 +312,6 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // One id per span plus the probe's headroom. Reserving *two* apiece was measured worse, not - // better: the allocating entry point sizes its buffer at `len/4`, which is about one id per - // span, so asking for two forced a reallocation on every call that would not otherwise have - // happened. Anything past this grows amortised, and a caller that wants no growth at all - // should reserve once and use `encode_generic_into`. output.reserve(spans.len() + MAX_INLINE_IDS); let mut capacity = output.capacity(); let mut cursor = output.len(); @@ -395,8 +324,6 @@ impl pipeline::Model for PipelineBPE { continue; } - // One capacity check per word, covering both the fold's single write and the probe's - // `MAX_INLINE_IDS` lanes. After it, writing that many past `cursor` is in bounds. if cursor + MAX_INLINE_IDS > capacity { // SAFETY: `cursor` counts what has been written so far. unsafe { output.set_len(cursor) }; @@ -404,12 +331,6 @@ impl pipeline::Model for PipelineBPE { capacity = output.capacity(); } - // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a - // word that is itself a foldable vocabulary entry in one probe, and those words never - // reach the cache. Probing the cache first would populate it with words the fold - // already serves for free, and the two paths would disagree about what it holds. - // The span is inside `chunk`, so everything up to the chunk's end is readable and a - // short word's key can be one unaligned masked load rather than a head/tail stitch. let (key, hash) = key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); if let Some(id) = self.fold_id_keyed(key, hash) { @@ -438,8 +359,6 @@ impl pipeline::Model for PipelineBPE { cursor += n; continue; } - // A hit the fast path could not serve: the probe already found the ids, so copy - // those rather than probing a second time. ProbeEmit::Hit(ids) => { // SAFETY: `cursor` counts what has been written so far. unsafe { output.set_len(cursor) }; @@ -457,7 +376,6 @@ impl pipeline::Model for PipelineBPE { unsafe { output.set_len(cursor) }; let start = output.len(); self.merge_word(sequence, symbols, queue); - // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), })); @@ -489,13 +407,6 @@ mod fold_tests { use crate::pipeline::PipelineTokenizer; /// The proven fold emits a vocabulary entry without merging, so it is only valid if the merge - /// loop would have produced that same entry. gpt2 does not declare `ignore_merges`, so here - /// the fold is on purely because the proof enabled it -- which makes it the config where a - /// wrong proof would show up. - /// - /// These strings mix words that are a single vocabulary entry (folded) with words that are - /// not (merged), and include the special token whose entry does NOT fold: `<|endoftext|>` - /// decomposes to seven tokens, and folding it would emit one. #[test] fn the_proven_fold_never_changes_the_ids() { let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); @@ -527,46 +438,4 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } - - /// `PipelineBPE::tokenize_spans` is an override of a trait method whose default is the - /// `tokenize_pipeline` loop, so the two can drift apart without anything failing to build -- - /// which is how it came to destructure a `BpeScratch` that no longer had those fields. - /// - /// The short strings above pass through the batch loop a handful of spans at a time. This one - /// gives it thousands in a single chunk, with the traffic that separates the two paths: - /// repeats (so the word cache both fills and hits), words the fold serves, words that must - /// merge, punctuation runs, multi-byte scripts, and a long unbroken run. - #[test] - fn the_batched_path_matches_the_reference() { - let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); - let pipe = PipelineTokenizer::try_from(&reference).unwrap(); - - let mut text = String::new(); - for i in 0..400 { - text.push_str(" the quick brown fox jumps over the lazy dog"); - text.push_str(" internationalisation unfortunately"); - text.push_str(" def foo(bar): return bar + 1"); - text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); - text.push_str(" 语言模型 ελληνικά"); - if i % 3 == 0 { - text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); - } - } - - let want: Vec = reference - .encode_fast(text.as_str(), false) - .unwrap() - .get_ids() - .to_vec(); - let got: Vec = pipe - .encode(text.as_str(), false) - .wait() - .unwrap() - .remove(0) - .iter() - .map(|t| t.id) - .collect(); - assert_eq!(want.len(), got.len(), "token count differs"); - assert_eq!(want, got, "the batched path changed the ids"); - } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index b8c2982cc..7e857bccd 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -1,54 +1,3 @@ -//! A table that remembers which token ids a word encodes to, so the tokenization -//! model (the expensive step) only encodes each word once. -//! -//! We determine the placement of a word in the cache table with a **hash** of the word: -//! - the bottom bits pick its **home slot**. A word can be cached in a 16-slot window around it -//! - the top byte is a **tag**, stored in a separate table ([`WordCache::quick_lookup`]) -//! -//! A lookup walks a 16 byte window in [`WordCache::quick_lookup`] to find a matching tag, if the tag -//! matches the slot's 128 bit key ([`LookupKey`]) confirms whether it's a match or not. -//! -//! On miss, we return where the cache should insert ([`WordCache::insert`]) the ids; -//! Either the slot already holding a stale copy of the word ([`WordCacheSlot::is_stale`]), -//! the first empty (0x00) slot in the window, or the home slot if the window is full. -//! -//! ```text -//! lookup "hat": tag A7, home slot 5 -//! -//! slot: 4 5 6 7 -//! ┌─────┬─────┬─────┬─────┬──── -//! tags │ C4 │ A7 │ 31 │ A7 │ ... one hash byte per slot -//! └─────┴─────┴─────┴─────┴──── -//! ┌─────┬─────┬─────┬─────┬──── -//! slots │"cat"│"the"│"sat"│"hat"│ ... the key and the ids, 32 bytes -//! └─────┴─────┴─────┴─────┴──── -//! ▲ ▲ -//! │ └ tag and key match: a hit, return the ids -//! └ same tag, wrong key: keep walking -//! ``` -//! -//! A slot ([`WordCacheSlot`]) keeps up to three ids inline. -//! Longer encodings go to one shared buffer ([`WordCache::spilled_buffer`]) and the slot holds offsets in that buffer. -//! -//! # Note -//! -//! A cache hit for a word of 15 bytes or shorter is guaranteed to return correct ids. -//! For longer words, the cache hit relies on equality of 127 bits of hash of the word's bytes. -//! Two long words can in principle share the same 127 bit hash (a collision) which could make the -//! cache return incorrect ids for one of them, even though the collision is extremely unlikely. -//! -//! # Where the ideas come from -//! -//! - [Swiss Tables] is where the tag row comes from: one byte of hash per slot, -//! checked before the slot itself is touched. -//! - [gigatoken] is a BPE tokenizer with a pre-token cache built from the same -//! parts: `u128` packed keys, 32-byte self-contained slots, ids inline. -//! - [huggingface/tokenizers#2234] is an open-addressed cache for this same encode -//! pipeline, arrived at in parallel. -//! -//! [Swiss Tables]: https://abseil.io/about/design/swisstables -//! [gigatoken]: https://github.com/marcelroed/gigatoken -//! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 use std::fmt::Debug; @@ -59,15 +8,8 @@ use crate::vocab::bucket_vocab_store::key_and_hash; #[cfg(test)] use crate::vocab::bucket_vocab_store::INLINE_KEY_BYTES; -// One hash pass, not two, and shared with the vocabulary. -// -// This used to hash a long word twice -- `PLACEMENT_HASHER` for its slot and `DISCRIMINANT_HASHER` -// for the other half of its key -- which is why long pretokens paid for the key. It now keys through -// `bucket_vocab_store::key_and_hash`, the same function the fold probes with, so a word that misses -// the fold and falls through to this table is hashed once for both rather than once each. /// How many ids a [`WordCacheSlot`] holds inline before it has to spill. A probe writes this -/// many lanes unconditionally, so it is also the headroom [`WordCache::probe_emit`] needs. pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) @@ -91,15 +33,11 @@ impl<'a> WordCache { const EMPTY: u8 = 0; /// How many slots a word's window spans, starting at its home slot. - /// 16 tags are 16 bytes: a whole window fits in one vector register (SIMD) - /// and in one cache line, so scanning it for a hit can be a couple of instructions - /// and a single memory read. const WINDOW_SIZE: usize = 16; pub fn new(num_slots: usize) -> Self { let next_pow2 = num_slots.next_power_of_two(); if next_pow2 != num_slots { - // todo: warn the user the capacity has been rounded up } let n: usize = next_pow2 + Self::WINDOW_SIZE; let spilled_budget = 16 * n; @@ -114,16 +52,12 @@ impl<'a> WordCache { } /// Looks up a word in the cache. - /// On [Lookup::Hit], returns the ids it encodes to. - /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { self.lookup_placed(make_lookup_key(word, self.placement_mask)) } /// [`Self::lookup`] for a caller that already has the word's key and hash from - /// [`key_and_hash`] -- the fold probes the vocabulary with the same pair, so sharing it means a - /// word that misses the fold and falls through to here is hashed once, not twice. #[inline] pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) @@ -132,18 +66,15 @@ impl<'a> WordCache { /// Probe and emit in one step: on an inline hit in the home slot the ids are written straight /// to `dst` and the count returned, so nothing goes back to the slot and nothing becomes a /// slice. - /// /// This is the shape the hot path wants. [`Self::lookup`] hands back a `&[u32]`, which means /// the caller re-reads the slot to build a fat pointer and then copies a run whose length it /// only learns at run time -- three trips over one 32-byte line that a single load already /// brought in. Here that line is read once, all [`MAX_INLINE_IDS`] lanes are stored /// unconditionally, and the caller advances its cursor by the count: no branch on the length, /// no second load, no slice. - /// /// Falls back to the full window walk for anything else. The table is sized well above its /// load, so a word's home slot is usually the one it was placed in and the walk is a few /// percent of words. - /// /// # Safety /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. `word` must not be empty -- /// an empty word keys to zero, which is also what an untouched slot holds. @@ -155,7 +86,6 @@ impl<'a> WordCache { } /// [`Self::probe_emit`] for a caller that already has the word's key and hash. - /// /// # Safety /// As [`Self::probe_emit`]: `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. #[inline] @@ -164,8 +94,6 @@ impl<'a> WordCache { // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is // `next_pow2 + WINDOW_SIZE` long, so the home slot is always in bounds. let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; - // An untouched slot holds `LookupKey(0)`, which no non-empty word can key to, so a key - // match here is a real hit -- the same 127-bit argument the window walk makes. if slot.key == placement.key && !slot.is_spilled() { // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. Lanes past `ids_len` are // dead: the caller advances its cursor by `ids_len` only, so the next word overwrites @@ -184,7 +112,6 @@ impl<'a> WordCache { } /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup`] so - /// [`Self::probe_emit`] can fall back to it without hashing the word a second time. #[inline] fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { @@ -197,11 +124,9 @@ impl<'a> WordCache { let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); for candidate in candidates { - // Must validate that a candidate is indeed a match let slot = &self.cached_words[candidate]; if slot.key == key { if slot.is_stale(self.spilled_generation) { - // The entry is stale: replace it with fresh ids return Lookup::Miss(InsertPlacement { index: candidate, key, @@ -213,7 +138,6 @@ impl<'a> WordCache { } } - // No match: it's a miss. The ids go in the window's first empty slot or to the home slot Lookup::Miss(InsertPlacement { index: first_empty.unwrap_or(home), key, @@ -222,8 +146,6 @@ impl<'a> WordCache { } /// Insert a new (word, ids) pair in the cache - /// - /// The [InsertPlacement] comes from [`Lookup::Miss`] pub fn insert(&mut self, placement: InsertPlacement, ids: impl ExactSizeIterator) { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; @@ -232,9 +154,7 @@ impl<'a> WordCache { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { - // Spilled buffer budget passed: we clear it self.spilled_buffer.clear(); - // Bump the generation to invalidate previous spilled slots self.spilled_generation = self.spilled_generation.wrapping_add(1); if self.spilled_generation == 0 { self.reset(); @@ -263,7 +183,6 @@ impl<'a> WordCache { } fn reset(&mut self) { - // todo: log the cache clear self.spilled_buffer.clear(); self.quick_lookup = vec![0; self.quick_lookup.len()].into_boxed_slice(); self.cached_words = @@ -309,13 +228,11 @@ pub struct WordCacheSlot { _pad: [u8; 3], } -// 32-byte size and alignment, so a read is always contained in a cache line const _: () = assert!(size_of::() == 32); const _: () = assert!(align_of::() == 32); impl WordCacheSlot { /// A sentinel value that discriminates an inline slot (ids are stored in the slot) from a spilled slot - /// (values are stored in a buffer) const SPILLED: u8 = 0xFF; pub fn new_spilled(key: LookupKey, ids_offsets: (usize, usize), generation: u32) -> Self { @@ -366,7 +283,6 @@ impl WordCacheSlot { } /// Convenience wrapper around [`WordCacheSlot`] to discriminate -/// whether it's a spilled or self-contained slot pub enum CacheSlotType<'a> { SelfContained(SelfContained<'a>), Spilled(Spilled<'a>), @@ -430,10 +346,6 @@ impl<'a> SelfContained<'a> { pub struct LookupKey(u64); /// The key, home slot and tag of a word. -/// -/// The key and its hash come from [`key_and_hash`], which is also what the vocabulary store keys on -/// -- one scheme, so a word that is probed in both tables can be hashed once. See -/// [`WordCache::lookup_keyed`]. #[inline] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { let (key, hash) = key_and_hash(word); @@ -446,14 +358,11 @@ fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlace key, index: (hash & placement_mask) as usize, tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), - // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } impl std::fmt::Debug for LookupKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // An inline key carries its length in the top byte and its bytes below; anything else is a - // hash and has nothing readable in it. let len = (self.0 >> 56) as usize; if len <= 7 { let bytes = self.0.to_le_bytes(); @@ -478,13 +387,10 @@ pub enum Lookup<'a> { } /// What [`WordCache::probe_emit`] found. `Wrote` is the fast path: the ids are already at the -/// caller's cursor and only the count comes back. pub enum ProbeEmit<'a> { /// An inline hit in the home slot. [`MAX_INLINE_IDS`] lanes were written at `dst`; this many - /// of them are live. Wrote(usize), /// A hit the fast path could not serve -- a spilled entry, or one placed off its home slot. - /// The ids were found, so the caller copies these rather than probing again. Hit(&'a [u32]), Miss(InsertPlacement), } @@ -512,7 +418,6 @@ impl Window { .simd_eq(i8x16::from([WordCache::EMPTY as i8; 16])) .to_bitmask() as u16; - // a trick to truncate matches to before the first empty let before_first_empty = !empty_bitmask & empty_bitmask.wrapping_sub(1); let candidates = SlotSet { @@ -566,7 +471,6 @@ mod tests { } /// The first `n` probe words whose home slot falls in `range`, for tests that - /// need to pick where in the table their words land. fn words_homed_in(cache: &WordCache, range: std::ops::Range, n: usize) -> Vec> { (0u32..) .map(|i| format!("w{i}").into_bytes()) @@ -580,7 +484,6 @@ mod tests { } /// Three ids per word, since fewer has its own test, and homes clear of the - /// table's end, which has its own test too. #[test] fn a_stored_word_is_found_again() { let mut cache = WordCache::new(1 << 8); @@ -602,9 +505,6 @@ mod tests { } /// Home slots stop at `placement_mask`; the table holds [`WordCache::WINDOW_SIZE`] - /// extra slots so a window starting on the last home does not run out of entries. - /// Two words homed there force the second one into the extra slots; both have to - /// round trip anyway. #[test] fn words_homed_on_the_last_slot_round_trip() { let mut cache = WordCache::new(1 << 2); @@ -623,9 +523,6 @@ mod tests { } /// Both sides of the fifteen-byte key boundary, crossed with both sides of the - /// three-id inline boundary. The words past fifteen bytes are the ones whose - /// stored key is a hash ([`LookupKey::new_hash`]); no shorter word exercises - /// that path. #[test] fn every_key_and_slot_shape_round_trips() { let mut cache = WordCache::new(1 << 8); @@ -654,20 +551,12 @@ mod tests { } /// Two distinct words must never share a key, or a hit returns the other word's - /// ids. The first pair differs only in the top bit of the last byte, which any - /// word ending in a multi-byte UTF-8 character has set; the second pair differs - /// only in a trailing zero byte, so only the length tells them apart. And no - /// inline key may carry the bit that marks a hashed one, whatever its bytes. #[test] fn packed_keys_are_unique_per_word() { let cache = WordCache::new(1 << 8); let key = |word: &[u8]| make_lookup_key(word, cache.placement_mask).key; - // Past the inline range these are hashes, so distinctness is the hash's job, not the - // packing's -- but they must still not collide. assert_ne!(key(b"aaaaaaaaaaaaaa\x7f"), key(b"aaaaaaaaaaaaaa\xff")); - // Inside it, the length is part of the key, so a trailing NUL cannot be lost. assert_ne!(key(b"abcd"), key(b"abcd\0")); - // Every distinct word within the inline range gets a distinct key, by construction. let mut seen = std::collections::HashSet::new(); for len in 1..=INLINE_KEY_BYTES { for b in 0..=255u8 { @@ -678,14 +567,10 @@ mod tests { } /// One window shape per row: the needle in various lanes, an empty slot in - /// the middle, at the edges, or absent. Candidates past the first empty - /// lane must not be reported (no entry can live there, since inserts always - /// fill the first empty lane) and the empty lane itself is the placement. #[test] fn the_scan_reports_matches_before_the_first_empty_and_the_empty_itself() { let offset = 3; let cases: &[(&[u8], u16, Option)] = &[ - // needle at lanes 0 and 2, empty at 3: lane 5's match is out of reach ( &[ 0xA7, 0x31, 0xA7, 0x00, 0x5F, 0xA7, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, @@ -694,7 +579,6 @@ mod tests { 0b101, Some(3), ), - // no empty slot: every match is reachable, no placement ( &[ 0xA7, 0x31, 0xA7, 0x22, 0x5F, 0xA7, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, @@ -703,7 +587,6 @@ mod tests { 0b1000000000100101, None, ), - // empty in lane 0: nothing is reachable ( &[ 0x00, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, @@ -712,7 +595,6 @@ mod tests { 0, Some(0), ), - // every lane matches, empty in the last lane ( &[ 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, @@ -742,9 +624,6 @@ mod tests { } /// The tag row marks a free slot with 0x00, so no live entry may carry that - /// tag: the scan stops looking at the first 0x00 it sees, and an insert - /// treats the slot as free space. Enough words that a tag built from a raw - /// hash byte would land on 0x00 hundreds of times. #[test] fn a_live_tag_is_never_the_empty_marker() { for i in 0..100_000u32 { @@ -754,11 +633,6 @@ mod tests { } /// Every inline length, with a different value in every byte position, so a - /// packing that drops, duplicates or misplaces a byte fails. The reference is the construction - /// the packing must be equivalent to: the bytes in a zeroed `u64`, the length in the top byte. - /// - /// Only up to [`INLINE_KEY_BYTES`]; past that the key is a hash and there are no bytes in it to - /// check. #[test] fn an_inline_key_is_the_words_bytes_with_the_length_on_top() { for len in 0..=INLINE_KEY_BYTES { @@ -775,7 +649,6 @@ mod tests { } /// A nonzero start, since every spill after the first has one and offsets that - /// only round trip from zero would still pass. #[test] fn a_spilled_words_offsets_round_trip() { let cached = WordCacheSlot::new_spilled(LookupKey::default(), (5, 9), 0); @@ -797,9 +670,6 @@ mod tests { } /// A tag is one byte of hash, so about one occupied slot in 256 carries the tag - /// of a word it does not hold. Too rare to hit by chance here, so the collision - /// is forged. The lookup must confirm the key and keep walking, not trust the - /// tag. #[test] fn a_tag_collision_is_confirmed_against_the_key() { let mut cache = WordCache::new(1 << 8); @@ -815,11 +685,8 @@ mod tests { } /// A word whose whole window is taken is still cached: it evicts its home slot. - /// The other fifteen words in the window keep their ids. #[test] fn a_full_window_evicts_only_the_home_slot() { - // A one-home table: placement_mask is 0, so every word homes at slot 0 and - // the sixteen slots from there are one shared window. let mut cache = WordCache::new(1); let words = words_homed_in(&cache, 0..1, WordCache::WINDOW_SIZE); for (i, word) in words.iter().enumerate() { @@ -828,7 +695,6 @@ mod tests { store(&mut cache, b"newcomer", &[999]); assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999][..])); - // Inserts fill the window in order, so the home slot holds the first word. assert_eq!(cache.lookup(&words[0]).hit(), None); for (i, word) in words.iter().enumerate().skip(1) { assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); @@ -836,9 +702,6 @@ mod tests { } /// The filler word brings a whole budget of ids on its own, so caching it - /// evicts whatever the buffer holds. The evicted word's slot keeps its tag - /// and key, but its ids are gone from the buffer: looking the word up - /// again has to be a miss, not a hit on the filler's ids. #[test] fn a_spilled_word_misses_after_an_evict() { let mut cache = WordCache::new(1 << 6); @@ -848,9 +711,6 @@ mod tests { } /// An evict drops a word's ids but leaves its slot behind. The miss the - /// word's lookup then returns must be a placement `insert` accepts, and - /// the word must round trip through it. New ids for the second insertion, - /// so a hit on the first ones cannot pass. #[test] fn an_evicted_word_round_trips_once_re_inserted() { let mut cache = WordCache::new(1 << 6); @@ -864,9 +724,6 @@ mod tests { } /// An evict leaves the word's slot behind, key and tag intact. The miss the - /// word's lookup then returns must place the fresh ids back into that slot, - /// not into an empty one: the window stays clean of stale copies, which is - /// what lets the walk stop at the first key match. #[test] fn a_miss_reuses_the_slot_of_its_stale_copy() { let mut cache = WordCache::new(1 << 6); @@ -883,8 +740,6 @@ mod tests { } /// Five evict cycles on the same word: it must end up holding exactly one - /// slot. A placement that preferred an empty slot would leave one more - /// stale copy behind per cycle. #[test] fn an_evicted_word_never_occupies_two_slots() { let mut cache = WordCache::new(1 << 6); @@ -902,8 +757,6 @@ mod tests { } /// A self-contained word keeps its ids in the slot, not in the buffer, so - /// an evict must not cost it its hit. Two filler words, because the first - /// fills an empty buffer exactly to the budget; the second overflows it. #[test] fn a_self_contained_word_still_hits_after_an_evict() { let mut cache = WordCache::new(1 << 6); @@ -914,8 +767,6 @@ mod tests { } /// The evict happens inside the insert that caches this word: the slot - /// must carry the generation the evict moved to, not the one the insert - /// started with, or the word would be stale the moment it is cached. #[test] fn the_evicting_insert_caches_its_own_word() { let mut cache = WordCache::new(1 << 6); @@ -928,9 +779,6 @@ mod tests { } /// The generation counter wraps back to zero after 2^32 evicts, where a - /// slot stamped long ago would look fresh again and hit on another word's - /// ids. The wrap therefore clears the whole table, self-contained slots - /// included; only the word the wrapping insert caches survives it. #[test] fn the_evict_that_wraps_the_generation_clears_the_table() { let mut cache = WordCache::new(1 << 6); @@ -965,8 +813,6 @@ mod tests { } /// Churn a table far too small for its input, mixing every key and slot shape, - /// and check the one hard promise: the cache may forget a word, it must never - /// answer with another word's ids. #[test] fn a_hit_never_returns_another_words_ids() { let mut cache = WordCache::new(64); @@ -979,8 +825,6 @@ mod tests { _ => format!("{}-{i}", "z".repeat(i % 40)), } .into_bytes(); - // No two words share an id (k stays below 16), so a wrong hit cannot - // return the right ids by luck. let ids: Vec = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); store(&mut cache, &word, &ids); expected.push((word, ids)); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 896c7ab34..188da525b 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -5,12 +5,6 @@ use ptr_hash::{FastPtrHash, PtrHashParams, hash::NoHash}; use std::fmt; /// Hashes a word key. Fixed seeds so a vocabulary always hashes identically. -/// -/// One pass. xxh3-128 was tried, to give a long key 63 bits of discrimination independent of the 64 -/// that place it; it cost 5.8 -> 7.0 ns per probe and ~10-25% on chinese and russian, whose pretokens -/// are mostly long, so it was dropped. A long key reuses its placement hash as its discriminant and -/// adds the length instead: a false hit needs a 64-bit collision at equal length (~2^-64 per query) -/// rather than being impossible as the old `memcmp` made it. static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, @@ -20,34 +14,16 @@ static KEY_HASHER: RandomState = RandomState::with_seeds( type Mphf = FastPtrHash; -// No hasher on the struct: both hashes below are fixed, so build and query agree without one -// having to be carried along to keep them consistent. /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be -/// emitted without running the merge loop. See `PipelineBPE::prove_fold`. -/// -/// Packed into the id rather than kept beside it, following the same convention as a packed merge -/// value, where the low bits are the product id and the high bits are a flag field (`SAFE_MASK`). -/// The lookup already loads this entry to verify the key, so the flag costs no memory, no extra -/// load, and no second structure to keep in step with the vocabulary. const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; /// Tokens up to this many bytes are their own key: the bytes fit beside the length in a `u64`. -/// -/// Seven, not fifteen, because that is what the corpus is. English averages 4.83 bytes per -/// pretoken and code 4.08, so a `u128` key was paying double width, a two-part head/tail read and a -/// 32-byte entry to describe words that fit in a register. pub(crate) const INLINE_KEY_BYTES: usize = 7; /// Mixes a short key into the well-distributed `u64` the MPHF wants. -/// -/// One multiply and one shift. An inline key *is* the token, so the compare is exact no matter how -/// the slot was chosen -- the hash only has to spread well enough for the MPHF to separate the keys, -/// and aHash's rounds, or splitmix64's second multiply, are wasted on that. Dropping the mixing -/// entirely does not work: packed short keys share their high bytes, and construction fails outright -/// with "indistinguishable hashes in bucket". #[inline(always)] fn mix(z: u64) -> u64 { let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); @@ -55,22 +31,18 @@ fn mix(z: u64) -> u64 { } /// The token's bytes as a fixed-width key, and the `u64` the MPHF is indexed by, from one pass. -/// /// * up to [`INLINE_KEY_BYTES`] bytes: the key *is* the bytes, with the length in the top byte, so /// the compare is proof and `"ab"` cannot collide with `"ab\0"`. /// * longer: the key is aHash of the bytes, which mixes the length in, and doubles as the placement /// hash. A false hit then needs a full 64-bit collision (~2^-64 per query) rather than being /// impossible as a `memcmp` against the byte arena made it. -/// /// [`crate::utils::word_cache`] keys through this very function, so a pretoken probed in both tables /// is hashed once for the pair. /// [`key_and_hash`] when the caller can guarantee `readable` bytes exist from the word's start. -/// /// A pretoken sits inside a chunk, so as long as it is not within 8 bytes of the chunk's end, its /// key can be one unaligned 8-byte load masked to the length -- instead of a head load, a tail load /// and a variable shift to stitch them. `pack` measured 1.6-2.1 ns/word, the largest single piece of /// the fold path, and the fold answers 88-94% of pretokens. -/// /// # Safety /// Reading is safe for any `readable >= 8`; the mask discards whatever came from past the word. #[inline] @@ -99,8 +71,6 @@ pub fn key_and_hash(word: &[u8]) -> (u64, u64) { let hash = KEY_HASHER.hash_one(word); return (hash, hash); } - // One unaligned load of the whole key range, masked to the length. Reading past the word is not - // allowed, so read the tail and shift: still register-only, no `memcpy`. let raw = if len >= 4 { let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u64; let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u64; @@ -118,16 +88,6 @@ pub fn key_and_hash(word: &[u8]) -> (u64, u64) { } /// One probe entry: a digest to confirm the slot, and the id to return. **8 bytes.** -/// -/// The probe is the encode path's most expensive step -- 5.08 ns per span, measured, which is an L2 -/// miss: the MPHF scatters a corpus's few thousand hot words across the whole entry table, so each -/// one lands on its own line. Halving the entry halves the lines the hot set occupies. 32 bytes -> -/// 16 -> 8 across this session, and 50257 entries is now 400 KB where it started at 1.6 MB. -/// -/// A 32-bit digest, not the full key. Perfect hashing already guarantees that an *in-vocabulary* -/// word reaches its own slot, so the stored value only has to reject an out-of-vocabulary query -- -/// about 6% of latin pretokens. A wrong id needs one of those to collide in 32 bits: ~2^-32 per -/// missing word, against the ~2^-64 the long-key path already accepts. #[derive(Clone, Copy, Debug, Default)] #[repr(C)] struct Entry { @@ -139,11 +99,6 @@ struct Entry { const _: () = assert!(size_of::() == 8); /// `KEY_MASK[len]` keeps the low `len` bytes; `LEN_TAG[len]` is the length in the top byte. -/// -/// Two tiny always-resident loads instead of a four-deep dependent ALU chain -/// (`len -> 8*len -> 64-x -> shift -> and`). Packing the key measured 1.94 ns per span in situ, more -/// than the hash, the MPHF lookup and the entry load put together, and that chain is why: the loads -/// below issue in parallel with the word's own load, where the shifts could not. static KEY_MASK: [u64; 8] = [ 0x0000_0000_0000_0000, 0x0000_0000_0000_00FF, @@ -166,21 +121,13 @@ static LEN_TAG: [u64; 8] = [ ]; /// The 32 bits an entry stores to reject an out-of-vocabulary query. Derived from the key by a -/// different multiply than the placement hash, so it is not a restatement of the slot. #[inline(always)] /// The 32 bits that confirm a slot really holds the queried token. -/// -/// Taken from the hash rather than recomputed from the key. `mix` already multiplies the key by this -/// crate's odd constant, and the old digest multiplied by the *same* constant a second time, so every -/// pretoken paid two 64-bit multiplies where one does. `hash` is `m ^ (m >> 29)` for that product, so -/// its top half is still a deterministic, well-spread function of the key -- which is all a digest -/// has to be. Build and query both go through here, so they cannot disagree. fn digest_of(hash: u64) -> u32 { (hash >> 32) as u32 } /// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose: only the reverse lookup -/// wants it, and keeping it in `Entry` made every probe drag 8 dead bytes through cache. #[derive(Clone, Copy, Debug, Default)] struct Span { start: u32, @@ -217,8 +164,6 @@ pub struct BucketVocabStore { /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the - /// MPHF's non-minimal slot range (with phantom padding slots), so its length is not the - /// token count. n: usize, } @@ -236,7 +181,6 @@ impl PartialEq for BucketVocabStore { if self.len() != other.len() { return false; } - // early exit as soon as there is a missmatch for id in 0..self.len() { if self.id_to_token(id as u32) != other.id_to_token(id as u32) { return false; @@ -256,16 +200,11 @@ impl BucketVocabStore { pub fn build(tokens: Vec<(Vec, u32)>) -> Self { let n = tokens.len(); - // 1. Pre-hash token bytes -> u64 keys using near perfect hash func. - // Via the packed key, so build and query fold the same fixed-width value. let keys: Vec = tokens .iter() .map(|(s, _)| key_and_hash(s.as_slice()).1) .collect(); - // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely - // (~n^2/2^65); if one ever fires, switch the key type to u128. The byte check below makes - // a collision a correct miss at query time, but it would drop a token at build, so guard. let mut seen = HashSet::with_capacity(n); for k in &keys { @@ -281,17 +220,11 @@ impl BucketVocabStore { } } - // 3. Build the (non-minimal) `FastPtrHash` via `PtrHashParams::default_fast()`; query with `.index()`. let params = PtrHashParams::default_fast(); let mphf = Mphf::new(&seen.into_iter().collect::>(), params); - // FastPtrHash is non-minimal: `index()` may return a slot up to `max_index()` (>= n), - // so `entries` must be sized to cover the whole slot range. Slots never written by the - // build loop stay as the default `Entry { len: 0, .. }` (phantom/padding slots), which - // enumeration/count paths filter out via `len > 0`. let n_slots = mphf.max_index(); - // 4. Place each token at its MPHF slot; build the slab and the id->slot reverse table. let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); @@ -329,7 +262,6 @@ impl BucketVocabStore { } pub fn new() -> Self { - // convenient to build empty edit later. let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), @@ -342,9 +274,6 @@ impl BucketVocabStore { } /// This function is the equivalent of `get` on a HashaMap, it return the id - /// corresponding to the key `q`. Since `mphf` always return a slot, we check - /// whether the token indexed by that slot actually match the query. We don't - /// care about collisions on query because of this! #[inline] pub fn get_bytes(&self, q: &[u8]) -> Option { if self.entries.is_empty() { @@ -353,14 +282,10 @@ impl BucketVocabStore { let (key, hash) = key_and_hash(q); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - // Digest equality confirms `q` really is the token at this slot: perfect hashing only - // guarantees a valid slot for in-vocab keys, so this is what rejects an out-of-vocabulary - // query. An unwritten padding slot holds key 0, which no token can pack to. (e.digest == digest_of(hash)).then_some(e.id & VOCAB_ID_MASK) } /// The id for `q`, together with whether that entry may be folded. One probe and one entry - /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { let (key, hash) = key_and_hash(q); @@ -368,18 +293,12 @@ impl BucketVocabStore { } /// The slot a hash lands in. Split out of the probe so a caller with many words can issue all - /// the pilot loads before it needs any of the answers -- see [`Self::entry_at`]. #[inline(always)] pub fn probe_slot(&self, hash: u64) -> usize { self.mphf.index(&hash) } /// The `(key, id)` at a slot, without deciding anything. - /// - /// A probe is a chain of two dependent loads -- pilot, then entry -- and at one word at a time - /// the whole chain is exposed latency. A caller holding N words can run `probe_slot` for all of - /// them, then `entry_at` for all of them, and the CPU has N independent misses outstanding - /// instead of one. Same loads, same table, N times the memory parallelism. #[inline(always)] pub fn entry_at(&self, slot: usize) -> (u32, u32) { let e = self.entries[slot]; @@ -394,9 +313,6 @@ impl BucketVocabStore { } /// [`Self::get_bytes_foldable`] for a caller that already has the word's key and hash. - /// - /// The word cache keys words exactly the same way, so a pretoken that misses the fold and then - /// goes to the cache would otherwise be hashed twice. This lets one pass serve both. #[inline] pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { let _ = key; @@ -436,7 +352,6 @@ impl BucketVocabStore { #[inline] pub fn id_to_token(&self, id: u32) -> Option { - // we are not sure its a valid utf8 so if not, adds replacement char self.id_to_token_bytes(id) .map(|b| String::from_utf8_lossy(b).into_owned()) } @@ -446,10 +361,6 @@ impl BucketVocabStore { } /// One past the highest id this vocabulary can hold. - /// - /// Ids are not dense: a config may leave gaps, so [`Self::len`] counts entries and is *not* an - /// id bound. Anything that walks ids has to bound itself by this and skip the holes, which - /// [`Self::id_to_token_bytes`] reports as `None`. pub fn id_space(&self) -> usize { self.id_to_slot.len() } @@ -459,12 +370,10 @@ impl BucketVocabStore { } pub fn content(&self) -> Vec<(String, u32)> { - // `spans` says which slots the build actually wrote: a padding slot keeps length 0. self.entries .iter() .zip(self.spans.iter()) .filter(|(_, sp)| sp.len > 0) - // Mask: the stored id carries FOLD_BIT, which must never escape this type. .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() @@ -481,7 +390,6 @@ impl BucketVocabStore { .iter() .zip(self.spans.iter()) .filter(|(_, sp)| sp.len > 0) - // Mask: the stored id carries FOLD_BIT, which must never escape this type. .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() @@ -555,7 +463,6 @@ mod tests { #[test] fn eq_matches_on_dense_content() { - // Models use dense ids (0..n); equality must reflect the token set on that range. let a = BucketVocabStore::build(vec![(b"x".to_vec(), 0), (b"y".to_vec(), 1)]); let b = BucketVocabStore::build(vec![(b"y".to_vec(), 1), (b"x".to_vec(), 0)]); let c = BucketVocabStore::build(vec![(b"x".to_vec(), 0), (b"z".to_vec(), 1)]); From bb13400e749f0ce6b5549b57620ad9878284718f Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 15:28:41 +0900 Subject: [PATCH 14/16] drop the port's own rationale comments, keep the ones that were already there The measured-dead-end essays and alternatives-tried notes this port added belong in the PR, not the source: 167 comment lines removed across word_cache, bucket_vocab_store and model. Every comment that existed in the base is preserved verbatim. The 33 base comment lines that no longer appear are the ones whose subject the port deleted -- `PLACEMENT_HASHER` and `DISCRIMINANT_HASHER` (gone with the u64 key), "the hasher is also stored on the struct" (the field is gone), `entries[slot] -> (offset, length, id)` (an entry is now `(digest, id)`), and `fold_id`'s doc (folded into `fold_id_keyed`). `SAFETY` comments and doctests are untouched. --- tokenizers/tk-encode/src/models/bpe/model.rs | 107 ++++++++++++- tokenizers/tk-encode/src/utils/word_cache.rs | 151 +++++++++++++++--- .../tk-encode/src/vocab/bucket_vocab_store.rs | 68 ++++---- 3 files changed, 263 insertions(+), 63 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 919544327..08d283f31 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -1,3 +1,6 @@ +//! The pipeline BPE model: its tables, how it is built from a [`BPE`], and how a pretokenized +//! sequence is turned into tokens. Conversion to symbols lives in `convert`; the merge engines +//! are `merge_multipass` and `merge_hot_cold_queue`. use crate::models::bpe::At; use crate::models::bpe::Error; use crate::models::bpe::convert::{AFFIX_BUF, Affixes}; @@ -15,9 +18,12 @@ const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; /// The gate, indexed by a word's first content byte: words no longer than their gate go to +/// multipass, longer ones to the hot/cold queue. fn build_byte_to_gate() -> [u16; 256] { let mut b2g = [GATE_MULTI; 256]; b2g[..0x80].fill(GATE_ASCII); + // Kept for a word that is *only* a delimiter (a run of spaces), where there is no content to + // classify. Words with content are indexed past their delimiter -- see [`content_start`]. for ws in *b" \t\n\r" { b2g[ws as usize] = GATE_MULTI; } @@ -25,16 +31,22 @@ fn build_byte_to_gate() -> [u16; 256] { } /// ByteLevel produces `" word"` or `"Ġword"`, Metaspace produces `"▁word"`. Indexing byte 0 +/// classifies the delimiter instead of the content. #[inline] fn content_start(bytes: &[u8]) -> usize { match bytes { + // Metaspace `▁` (U+2581). [0xE2, 0x96, 0x81, rest @ ..] if !rest.is_empty() => 3, + // ByteLevel `Ġ` (U+0120) -- the byte-level spelling of a leading space. [0xC4, 0xA0, rest @ ..] if !rest.is_empty() => 2, + // A literal leading space, which a ByteLevel pre-tokenizer also hands over. [ws, rest @ ..] if ws.is_ascii_whitespace() && !rest.is_empty() => 1, _ => 0, } } +// The fused cache probe stores ids straight at a `*mut u32` pointing into the `Vec` +// the caller is filling. That is only sound while a token is layout-identical to its id. const _: () = assert!(size_of::() == size_of::()); const _: () = assert!(align_of::() == align_of::()); @@ -48,6 +60,7 @@ pub struct PipelineBPE { cache_capacity: Option, } +// A `PipelineBPE` holds exactly one `Atoms`, so `Chars`' 1 KB byte-fallback table costs nothing. #[allow(clippy::large_enum_variant)] pub(super) enum Atoms { /// The atoms are the 256 bytes; the symbol for each lives in `BpeTables::byte_internal`. @@ -61,16 +74,23 @@ pub(super) enum Atoms { impl PipelineBPE { /// True when this model was built with `with_byte_level`, which means + /// [`byte_level::transform_vocab`] already turned every vocabulary entry into its + /// **decoded raw bytes** at load time. Decoding is then a concatenation, and running a + /// `ByteLevel` decoder over these entries would decode a second time. pub(crate) fn is_byte_level(&self) -> bool { matches!(self.atoms, Atoms::Bytes) } /// A token's bytes, borrowed from the vocab store's slab. For a byte-level model these are + /// the decoded bytes (see [`Self::is_byte_level`]) and a single entry is not necessarily + /// valid UTF-8 on its own -- only the concatenation of a whole id sequence usually is. pub(crate) fn id_to_token_bytes(&self, id: u32) -> Option<&[u8]> { self.vocab.id_to_token_bytes(id) } /// A token as a `String`, for the decoder-chain route. Only meaningful when the entries are + /// the token strings as written, i.e. when [`Self::is_byte_level`] is false; a byte-level + /// model decodes through [`Self::id_to_token_bytes`] instead. pub(crate) fn id_to_token(&self, id: u32) -> Option { self.vocab.id_to_token(id) } @@ -91,6 +111,7 @@ impl PipelineBPE { cache, .. } = model; + // A capacity of zero means "no cache"; anything else sizes the per-scratch table. let cache_capacity = cache.map(|cache| cache.capacity).filter(|&c| c > 0); let prefix = continuing_subword_prefix.unwrap_or_default(); let suffix = end_of_word_suffix.unwrap_or_default(); @@ -103,6 +124,7 @@ impl PipelineBPE { merges, with_byte_level, ); + // the symbol stream is internal ids, mapped back through `unmap` at the very end let to_internal = |external: u32| -> Option { external_to_internal .get(external as usize) @@ -112,6 +134,7 @@ impl PipelineBPE { let (vocab, atoms) = if with_byte_level { let mut vocab = BucketVocabStore::build(vocab.byte_content()); vocab = byte_level::transform_vocab(vocab); + // every byte has to be an atom, or a word containing it could not be encoded at all for b in 0u8..=255 { vocab .get_bytes(&[b]) @@ -165,6 +188,13 @@ impl PipelineBPE { vocab, byte_to_gate: build_byte_to_gate(), }; + // Every entry carries a foldable bit, so the encode path is one probe and one bit test + // with no policy left in it. The policy is decided here, once: a config that declares + // `ignore_merges` asks for every hit to fold, so every entry gets the bit; otherwise only + // the entries that prove they reduce to themselves earn it. + // + // Two phases because the proof runs the merge engine, which borrows `built`: work out the + // answers first, then set the bit on each entry that earned it. let proven = if ignore_merges { vec![true; built.vocab.id_space()] } else { @@ -179,7 +209,12 @@ impl PipelineBPE { } /// One bit per vocabulary id: can a pretoken equal to this entry be emitted as this entry, + /// without running the merge loop? + /// + /// We replace the old "ignore_merges" with something that actually ignores whether or not the flag was set. fn prove_fold(&self) -> Vec { + // The id space, not the entry count: ids may be sparse, and bounding the walk by + // `vocab.len()` would leave every entry above it unproven. let len = self.vocab.id_space(); let mut proven = vec![false; len]; let mut symbols = Vec::with_capacity(64); @@ -188,10 +223,13 @@ impl PipelineBPE { let Some(bytes) = self.vocab.id_to_token_bytes(id) else { continue; }; + // An entry that is not valid UTF-8 can never equal a pretoken, which is always a + // `&str` slice, so it can never be folded and needs no proof. let Ok(text) = std::str::from_utf8(bytes) else { continue; }; let foldable = if text.chars().count() <= 1 { + // A single atom has no pair to merge and is trivially its own encoding. true } else { self.merge_word(text, &mut symbols, &mut scratch); @@ -202,14 +240,19 @@ impl PipelineBPE { proven } - /// The id to emit for a pretoken without merging, when the whole word is a vocabulary entry #[inline(always)] fn fold_id_keyed(&self, key: u64, hash: u64) -> Option { + // One probe; the foldable bit is part of the id that probe already returned. Which entries let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } /// Converts a word to symbols and merges it. The gate, indexed by the word's first *content* + /// byte (past any delimiter the pre-tokenizer prepended -- see [`content_start`]), says + /// which engine gets it: short words go to multipass, longer ones to the hot/cold queue. + /// `symbols` is the caller's reusable symbol buffer -- it lives in the scratch so that a word + /// costs no allocation. On return it holds the merged word as internal ids, which the caller + /// maps to external ids through `unmap`. pub(super) fn merge_word( &self, sequence: &str, @@ -217,9 +260,11 @@ impl PipelineBPE { queue_scratch: &mut QueueScratch, ) { let bytes = sequence.as_bytes(); + // Classify on the first content byte, not on the delimiter the pre-tokenizer prepended. let gate: u16 = self.byte_to_gate[bytes[content_start(bytes)] as usize]; if sequence.len() > gate as usize { + // conversion writes the entries and cold keys directly: no intermediate symbol array self.convert_queue( sequence, symbols, @@ -235,12 +280,14 @@ impl PipelineBPE { } /// Per-thread scratch for BPE. Every buffer here is cleared, never reallocated, so tokenizing a +/// sequence does not allocate. pub struct BpeScratch { /// Symbols of the word being merged. pub(crate) symbols: Vec, /// Entry arena and the two queue tiers. pub(crate) queue: QueueScratch, /// Words already seen, so a repeat costs a probe instead of a merge. It lives in the scratch + /// so it outlives the encode call that fills it -- otherwise it would never see a word twice. pub(crate) word_cache: Option, } @@ -272,6 +319,7 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; + // A word seen before costs a probe instead of a merge. let insert_at = if let Some(cache) = word_cache.as_mut() { match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { @@ -286,6 +334,7 @@ impl pipeline::Model for PipelineBPE { let start = output.len(); self.merge_word(sequence, symbols, queue); + // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), })); @@ -299,6 +348,11 @@ impl pipeline::Model for PipelineBPE { } /// Every pre-token of a chunk in one call. + /// + /// Same work per word as [`Self::tokenize_pipeline`]; what changes is what is *not* repeated. + /// The scratch is destructured once instead of once per word, the output is grown once for the + /// whole batch instead of being capacity-checked on every push, and the virtual call, the + /// slice and the `Result` happen once per chunk rather than once per pre-token. fn tokenize_spans( &self, chunk: &str, @@ -331,6 +385,10 @@ impl pipeline::Model for PipelineBPE { capacity = output.capacity(); } + // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a + // word that is itself a foldable vocabulary entry in one probe, and those words never + // reach the cache. Probing the cache first would populate it with words the fold + // already serves for free, and the two paths would disagree about what it holds. let (key, hash) = key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); if let Some(id) = self.fold_id_keyed(key, hash) { @@ -342,11 +400,7 @@ impl pipeline::Model for PipelineBPE { let mut placement = None; if let Some(cache) = word_cache.as_mut() { - // The probe writes the ids at the cursor itself, so a hit is one load of the slot - // and one unconditional store of its lanes -- the ids never become a slice and the - // line is never read twice. // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and - // `PipelineToken` is layout-identical to `u32` (asserted at the top of this file). let found = unsafe { cache.probe_emit_keyed( key, @@ -372,10 +426,10 @@ impl pipeline::Model for PipelineBPE { } // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` - // through its normal API, so its length has to be true again first. unsafe { output.set_len(cursor) }; let start = output.len(); self.merge_word(sequence, symbols, queue); + // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids output.extend(symbols.iter().map(|&symbol| PipelineToken { id: self.tables.unmap.at(symbol as usize), })); @@ -407,6 +461,13 @@ mod fold_tests { use crate::pipeline::PipelineTokenizer; /// The proven fold emits a vocabulary entry without merging, so it is only valid if the merge + /// loop would have produced that same entry. gpt2 does not declare `ignore_merges`, so here + /// the fold is on purely because the proof enabled it -- which makes it the config where a + /// wrong proof would show up. + /// + /// These strings mix words that are a single vocabulary entry (folded) with words that are + /// not (merged), and include the special token whose entry does NOT fold: `<|endoftext|>` + /// decomposes to seven tokens, and folding it would emit one. #[test] fn the_proven_fold_never_changes_the_ids() { let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); @@ -438,4 +499,38 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } + + #[test] + fn the_batched_path_matches_the_reference() { + let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); + let pipe = PipelineTokenizer::try_from(&reference).unwrap(); + + let mut text = String::new(); + for i in 0..400 { + text.push_str(" the quick brown fox jumps over the lazy dog"); + text.push_str(" internationalisation unfortunately"); + text.push_str(" def foo(bar): return bar + 1"); + text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); + text.push_str(" 语言模型 ελληνικά"); + if i % 3 == 0 { + text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); + } + } + + let want: Vec = reference + .encode_fast(text.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipe + .encode(text.as_str(), false) + .wait() + .unwrap() + .remove(0) + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want.len(), got.len(), "token count differs"); + assert_eq!(want, got, "the batched path changed the ids"); + } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 7e857bccd..9e0d6c02c 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -1,3 +1,54 @@ +//! A table that remembers which token ids a word encodes to, so the tokenization +//! model (the expensive step) only encodes each word once. +//! +//! We determine the placement of a word in the cache table with a **hash** of the word: +//! - the bottom bits pick its **home slot**. A word can be cached in a 16-slot window around it +//! - the top byte is a **tag**, stored in a separate table ([`WordCache::quick_lookup`]) +//! +//! A lookup walks a 16 byte window in [`WordCache::quick_lookup`] to find a matching tag, if the tag +//! matches the slot's 128 bit key ([`LookupKey`]) confirms whether it's a match or not. +//! +//! On miss, we return where the cache should insert ([`WordCache::insert`]) the ids; +//! Either the slot already holding a stale copy of the word ([`WordCacheSlot::is_stale`]), +//! the first empty (0x00) slot in the window, or the home slot if the window is full. +//! +//! ```text +//! lookup "hat": tag A7, home slot 5 +//! +//! slot: 4 5 6 7 +//! ┌─────┬─────┬─────┬─────┬──── +//! tags │ C4 │ A7 │ 31 │ A7 │ ... one hash byte per slot +//! └─────┴─────┴─────┴─────┴──── +//! ┌─────┬─────┬─────┬─────┬──── +//! slots │"cat"│"the"│"sat"│"hat"│ ... the key and the ids, 32 bytes +//! └─────┴─────┴─────┴─────┴──── +//! ▲ ▲ +//! │ └ tag and key match: a hit, return the ids +//! └ same tag, wrong key: keep walking +//! ``` +//! +//! A slot ([`WordCacheSlot`]) keeps up to three ids inline. +//! Longer encodings go to one shared buffer ([`WordCache::spilled_buffer`]) and the slot holds offsets in that buffer. +//! +//! # Note +//! +//! A cache hit for a word of 15 bytes or shorter is guaranteed to return correct ids. +//! For longer words, the cache hit relies on equality of 127 bits of hash of the word's bytes. +//! Two long words can in principle share the same 127 bit hash (a collision) which could make the +//! cache return incorrect ids for one of them, even though the collision is extremely unlikely. +//! +//! # Where the ideas come from +//! +//! - [Swiss Tables] is where the tag row comes from: one byte of hash per slot, +//! checked before the slot itself is touched. +//! - [gigatoken] is a BPE tokenizer with a pre-token cache built from the same +//! parts: `u128` packed keys, 32-byte self-contained slots, ids inline. +//! - [huggingface/tokenizers#2234] is an open-addressed cache for this same encode +//! pipeline, arrived at in parallel. +//! +//! [Swiss Tables]: https://abseil.io/about/design/swisstables +//! [gigatoken]: https://github.com/marcelroed/gigatoken +//! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 use std::fmt::Debug; @@ -9,7 +60,6 @@ use crate::vocab::bucket_vocab_store::key_and_hash; use crate::vocab::bucket_vocab_store::INLINE_KEY_BYTES; -/// How many ids a [`WordCacheSlot`] holds inline before it has to spill. A probe writes this pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) @@ -33,11 +83,15 @@ impl<'a> WordCache { const EMPTY: u8 = 0; /// How many slots a word's window spans, starting at its home slot. + /// 16 tags are 16 bytes: a whole window fits in one vector register (SIMD) + /// and in one cache line, so scanning it for a hit can be a couple of instructions + /// and a single memory read. const WINDOW_SIZE: usize = 16; pub fn new(num_slots: usize) -> Self { let next_pow2 = num_slots.next_power_of_two(); if next_pow2 != num_slots { + // todo: warn the user the capacity has been rounded up } let n: usize = next_pow2 + Self::WINDOW_SIZE; let spilled_budget = 16 * n; @@ -52,32 +106,22 @@ impl<'a> WordCache { } /// Looks up a word in the cache. + /// On [Lookup::Hit], returns the ids it encodes to. + /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { self.lookup_placed(make_lookup_key(word, self.placement_mask)) } - /// [`Self::lookup`] for a caller that already has the word's key and hash from #[inline] pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) } - /// Probe and emit in one step: on an inline hit in the home slot the ids are written straight - /// to `dst` and the count returned, so nothing goes back to the slot and nothing becomes a - /// slice. - /// This is the shape the hot path wants. [`Self::lookup`] hands back a `&[u32]`, which means - /// the caller re-reads the slot to build a fat pointer and then copies a run whose length it - /// only learns at run time -- three trips over one 32-byte line that a single load already - /// brought in. Here that line is read once, all [`MAX_INLINE_IDS`] lanes are stored - /// unconditionally, and the caller advances its cursor by the count: no branch on the length, - /// no second load, no slice. - /// Falls back to the full window walk for anything else. The table is sized well above its - /// load, so a word's home slot is usually the one it was placed in and the walk is a few - /// percent of words. + /// + /// + /// /// # Safety - /// `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. `word` must not be empty -- - /// an empty word keys to zero, which is also what an untouched slot holds. #[inline] pub unsafe fn probe_emit(&'a self, word: &[u8], dst: *mut u32) -> ProbeEmit<'a> { debug_assert!(!word.is_empty(), "probe_emit needs a non-empty word"); @@ -85,19 +129,15 @@ impl<'a> WordCache { unsafe { self.probe_emit_keyed(key, hash, dst) } } - /// [`Self::probe_emit`] for a caller that already has the word's key and hash. + /// /// # Safety - /// As [`Self::probe_emit`]: `dst` must have room for [`MAX_INLINE_IDS`] `u32` writes. #[inline] pub unsafe fn probe_emit_keyed(&'a self, key: u64, hash: u64, dst: *mut u32) -> ProbeEmit<'a> { let placement = placement_from(LookupKey(key), hash, self.placement_mask); // SAFETY: `index` is masked with `placement_mask` (`next_pow2 - 1`), and the table is - // `next_pow2 + WINDOW_SIZE` long, so the home slot is always in bounds. let slot = unsafe { *self.cached_words.as_ptr().add(placement.index) }; if slot.key == placement.key && !slot.is_spilled() { // SAFETY: the caller guarantees room for `MAX_INLINE_IDS`. Lanes past `ids_len` are - // dead: the caller advances its cursor by `ids_len` only, so the next word overwrites - // them or the final `set_len` cuts them off. unsafe { for lane in 0..MAX_INLINE_IDS { dst.add(lane).write(slot.payload[lane]); @@ -111,7 +151,6 @@ impl<'a> WordCache { } } - /// The window walk, once a word has been keyed and placed. Split out of [`Self::lookup`] so #[inline] fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { @@ -124,9 +163,11 @@ impl<'a> WordCache { let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); for candidate in candidates { + // Must validate that a candidate is indeed a match let slot = &self.cached_words[candidate]; if slot.key == key { if slot.is_stale(self.spilled_generation) { + // The entry is stale: replace it with fresh ids return Lookup::Miss(InsertPlacement { index: candidate, key, @@ -138,6 +179,7 @@ impl<'a> WordCache { } } + // No match: it's a miss. The ids go in the window's first empty slot or to the home slot Lookup::Miss(InsertPlacement { index: first_empty.unwrap_or(home), key, @@ -146,6 +188,8 @@ impl<'a> WordCache { } /// Insert a new (word, ids) pair in the cache + /// + /// The [InsertPlacement] comes from [`Lookup::Miss`] pub fn insert(&mut self, placement: InsertPlacement, ids: impl ExactSizeIterator) { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; @@ -154,7 +198,9 @@ impl<'a> WordCache { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { + // Spilled buffer budget passed: we clear it self.spilled_buffer.clear(); + // Bump the generation to invalidate previous spilled slots self.spilled_generation = self.spilled_generation.wrapping_add(1); if self.spilled_generation == 0 { self.reset(); @@ -183,6 +229,7 @@ impl<'a> WordCache { } fn reset(&mut self) { + // todo: log the cache clear self.spilled_buffer.clear(); self.quick_lookup = vec![0; self.quick_lookup.len()].into_boxed_slice(); self.cached_words = @@ -228,11 +275,13 @@ pub struct WordCacheSlot { _pad: [u8; 3], } +// 32-byte size and alignment, so a read is always contained in a cache line const _: () = assert!(size_of::() == 32); const _: () = assert!(align_of::() == 32); impl WordCacheSlot { /// A sentinel value that discriminates an inline slot (ids are stored in the slot) from a spilled slot + /// (values are stored in a buffer) const SPILLED: u8 = 0xFF; pub fn new_spilled(key: LookupKey, ids_offsets: (usize, usize), generation: u32) -> Self { @@ -283,6 +332,7 @@ impl WordCacheSlot { } /// Convenience wrapper around [`WordCacheSlot`] to discriminate +/// whether it's a spilled or self-contained slot pub enum CacheSlotType<'a> { SelfContained(SelfContained<'a>), Spilled(Spilled<'a>), @@ -358,6 +408,7 @@ fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlace key, index: (hash & placement_mask) as usize, tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), + // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } @@ -386,11 +437,8 @@ pub enum Lookup<'a> { Miss(InsertPlacement), } -/// What [`WordCache::probe_emit`] found. `Wrote` is the fast path: the ids are already at the pub enum ProbeEmit<'a> { - /// An inline hit in the home slot. [`MAX_INLINE_IDS`] lanes were written at `dst`; this many Wrote(usize), - /// A hit the fast path could not serve -- a spilled entry, or one placed off its home slot. Hit(&'a [u32]), Miss(InsertPlacement), } @@ -418,6 +466,7 @@ impl Window { .simd_eq(i8x16::from([WordCache::EMPTY as i8; 16])) .to_bitmask() as u16; + // a trick to truncate matches to before the first empty let before_first_empty = !empty_bitmask & empty_bitmask.wrapping_sub(1); let candidates = SlotSet { @@ -471,6 +520,7 @@ mod tests { } /// The first `n` probe words whose home slot falls in `range`, for tests that + /// need to pick where in the table their words land. fn words_homed_in(cache: &WordCache, range: std::ops::Range, n: usize) -> Vec> { (0u32..) .map(|i| format!("w{i}").into_bytes()) @@ -484,6 +534,7 @@ mod tests { } /// Three ids per word, since fewer has its own test, and homes clear of the + /// table's end, which has its own test too. #[test] fn a_stored_word_is_found_again() { let mut cache = WordCache::new(1 << 8); @@ -505,6 +556,9 @@ mod tests { } /// Home slots stop at `placement_mask`; the table holds [`WordCache::WINDOW_SIZE`] + /// extra slots so a window starting on the last home does not run out of entries. + /// Two words homed there force the second one into the extra slots; both have to + /// round trip anyway. #[test] fn words_homed_on_the_last_slot_round_trip() { let mut cache = WordCache::new(1 << 2); @@ -523,6 +577,9 @@ mod tests { } /// Both sides of the fifteen-byte key boundary, crossed with both sides of the + /// three-id inline boundary. The words past fifteen bytes are the ones whose + /// stored key is a hash ([`LookupKey::new_hash`]); no shorter word exercises + /// that path. #[test] fn every_key_and_slot_shape_round_trips() { let mut cache = WordCache::new(1 << 8); @@ -551,6 +608,10 @@ mod tests { } /// Two distinct words must never share a key, or a hit returns the other word's + /// ids. The first pair differs only in the top bit of the last byte, which any + /// word ending in a multi-byte UTF-8 character has set; the second pair differs + /// only in a trailing zero byte, so only the length tells them apart. And no + /// inline key may carry the bit that marks a hashed one, whatever its bytes. #[test] fn packed_keys_are_unique_per_word() { let cache = WordCache::new(1 << 8); @@ -567,10 +628,14 @@ mod tests { } /// One window shape per row: the needle in various lanes, an empty slot in + /// the middle, at the edges, or absent. Candidates past the first empty + /// lane must not be reported (no entry can live there, since inserts always + /// fill the first empty lane) and the empty lane itself is the placement. #[test] fn the_scan_reports_matches_before_the_first_empty_and_the_empty_itself() { let offset = 3; let cases: &[(&[u8], u16, Option)] = &[ + // needle at lanes 0 and 2, empty at 3: lane 5's match is out of reach ( &[ 0xA7, 0x31, 0xA7, 0x00, 0x5F, 0xA7, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, @@ -579,6 +644,7 @@ mod tests { 0b101, Some(3), ), + // no empty slot: every match is reachable, no placement ( &[ 0xA7, 0x31, 0xA7, 0x22, 0x5F, 0xA7, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, @@ -587,6 +653,7 @@ mod tests { 0b1000000000100101, None, ), + // empty in lane 0: nothing is reachable ( &[ 0x00, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, @@ -595,6 +662,7 @@ mod tests { 0, Some(0), ), + // every lane matches, empty in the last lane ( &[ 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, 0xA7, @@ -624,6 +692,9 @@ mod tests { } /// The tag row marks a free slot with 0x00, so no live entry may carry that + /// tag: the scan stops looking at the first 0x00 it sees, and an insert + /// treats the slot as free space. Enough words that a tag built from a raw + /// hash byte would land on 0x00 hundreds of times. #[test] fn a_live_tag_is_never_the_empty_marker() { for i in 0..100_000u32 { @@ -649,6 +720,7 @@ mod tests { } /// A nonzero start, since every spill after the first has one and offsets that + /// only round trip from zero would still pass. #[test] fn a_spilled_words_offsets_round_trip() { let cached = WordCacheSlot::new_spilled(LookupKey::default(), (5, 9), 0); @@ -670,6 +742,9 @@ mod tests { } /// A tag is one byte of hash, so about one occupied slot in 256 carries the tag + /// of a word it does not hold. Too rare to hit by chance here, so the collision + /// is forged. The lookup must confirm the key and keep walking, not trust the + /// tag. #[test] fn a_tag_collision_is_confirmed_against_the_key() { let mut cache = WordCache::new(1 << 8); @@ -685,8 +760,11 @@ mod tests { } /// A word whose whole window is taken is still cached: it evicts its home slot. + /// The other fifteen words in the window keep their ids. #[test] fn a_full_window_evicts_only_the_home_slot() { + // A one-home table: placement_mask is 0, so every word homes at slot 0 and + // the sixteen slots from there are one shared window. let mut cache = WordCache::new(1); let words = words_homed_in(&cache, 0..1, WordCache::WINDOW_SIZE); for (i, word) in words.iter().enumerate() { @@ -695,6 +773,7 @@ mod tests { store(&mut cache, b"newcomer", &[999]); assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999][..])); + // Inserts fill the window in order, so the home slot holds the first word. assert_eq!(cache.lookup(&words[0]).hit(), None); for (i, word) in words.iter().enumerate().skip(1) { assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); @@ -702,6 +781,9 @@ mod tests { } /// The filler word brings a whole budget of ids on its own, so caching it + /// evicts whatever the buffer holds. The evicted word's slot keeps its tag + /// and key, but its ids are gone from the buffer: looking the word up + /// again has to be a miss, not a hit on the filler's ids. #[test] fn a_spilled_word_misses_after_an_evict() { let mut cache = WordCache::new(1 << 6); @@ -711,6 +793,9 @@ mod tests { } /// An evict drops a word's ids but leaves its slot behind. The miss the + /// word's lookup then returns must be a placement `insert` accepts, and + /// the word must round trip through it. New ids for the second insertion, + /// so a hit on the first ones cannot pass. #[test] fn an_evicted_word_round_trips_once_re_inserted() { let mut cache = WordCache::new(1 << 6); @@ -724,6 +809,9 @@ mod tests { } /// An evict leaves the word's slot behind, key and tag intact. The miss the + /// word's lookup then returns must place the fresh ids back into that slot, + /// not into an empty one: the window stays clean of stale copies, which is + /// what lets the walk stop at the first key match. #[test] fn a_miss_reuses_the_slot_of_its_stale_copy() { let mut cache = WordCache::new(1 << 6); @@ -740,6 +828,8 @@ mod tests { } /// Five evict cycles on the same word: it must end up holding exactly one + /// slot. A placement that preferred an empty slot would leave one more + /// stale copy behind per cycle. #[test] fn an_evicted_word_never_occupies_two_slots() { let mut cache = WordCache::new(1 << 6); @@ -757,6 +847,8 @@ mod tests { } /// A self-contained word keeps its ids in the slot, not in the buffer, so + /// an evict must not cost it its hit. Two filler words, because the first + /// fills an empty buffer exactly to the budget; the second overflows it. #[test] fn a_self_contained_word_still_hits_after_an_evict() { let mut cache = WordCache::new(1 << 6); @@ -767,6 +859,8 @@ mod tests { } /// The evict happens inside the insert that caches this word: the slot + /// must carry the generation the evict moved to, not the one the insert + /// started with, or the word would be stale the moment it is cached. #[test] fn the_evicting_insert_caches_its_own_word() { let mut cache = WordCache::new(1 << 6); @@ -779,6 +873,9 @@ mod tests { } /// The generation counter wraps back to zero after 2^32 evicts, where a + /// slot stamped long ago would look fresh again and hit on another word's + /// ids. The wrap therefore clears the whole table, self-contained slots + /// included; only the word the wrapping insert caches survives it. #[test] fn the_evict_that_wraps_the_generation_clears_the_table() { let mut cache = WordCache::new(1 << 6); @@ -813,6 +910,8 @@ mod tests { } /// Churn a table far too small for its input, mixing every key and slot shape, + /// and check the one hard promise: the cache may forget a word, it must never + /// answer with another word's ids. #[test] fn a_hit_never_returns_another_words_ids() { let mut cache = WordCache::new(64); @@ -825,6 +924,8 @@ mod tests { _ => format!("{}-{i}", "z".repeat(i % 40)), } .into_bytes(); + // No two words share an id (k stays below 16), so a wrong hit cannot + // return the right ids by luck. let ids: Vec = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); store(&mut cache, &word, &ids); expected.push((word, ids)); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 188da525b..35f76e369 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,7 +4,6 @@ use ahash::RandomState; use ptr_hash::{FastPtrHash, PtrHashParams, hash::NoHash}; use std::fmt; -/// Hashes a word key. Fixed seeds so a vocabulary always hashes identically. static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, @@ -16,35 +15,29 @@ type Mphf = FastPtrHash; /// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be +/// emitted without running the merge loop. See `PipelineBPE::prove_fold`. +/// +/// Packed into the id rather than kept beside it, following the same convention as a packed merge +/// value, where the low bits are the product id and the high bits are a flag field (`SAFE_MASK`). +/// The lookup already loads this entry to verify the key, so the flag costs no memory, no extra +/// load, and no second structure to keep in step with the vocabulary. const FOLD_BIT: u32 = 1 << 31; /// The id half. 2^31 ids is far past any vocabulary. const VOCAB_ID_MASK: u32 = FOLD_BIT - 1; -/// Tokens up to this many bytes are their own key: the bytes fit beside the length in a `u64`. pub(crate) const INLINE_KEY_BYTES: usize = 7; -/// Mixes a short key into the well-distributed `u64` the MPHF wants. #[inline(always)] fn mix(z: u64) -> u64 { let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); z ^ (z >> 29) } -/// The token's bytes as a fixed-width key, and the `u64` the MPHF is indexed by, from one pass. -/// * up to [`INLINE_KEY_BYTES`] bytes: the key *is* the bytes, with the length in the top byte, so -/// the compare is proof and `"ab"` cannot collide with `"ab\0"`. -/// * longer: the key is aHash of the bytes, which mixes the length in, and doubles as the placement -/// hash. A false hit then needs a full 64-bit collision (~2^-64 per query) rather than being -/// impossible as a `memcmp` against the byte arena made it. -/// [`crate::utils::word_cache`] keys through this very function, so a pretoken probed in both tables -/// is hashed once for the pair. -/// [`key_and_hash`] when the caller can guarantee `readable` bytes exist from the word's start. -/// A pretoken sits inside a chunk, so as long as it is not within 8 bytes of the chunk's end, its -/// key can be one unaligned 8-byte load masked to the length -- instead of a head load, a tail load -/// and a variable shift to stitch them. `pack` measured 1.6-2.1 ns/word, the largest single piece of -/// the fold path, and the fold answers 88-94% of pretokens. +/// +/// +/// +/// /// # Safety -/// Reading is safe for any `readable >= 8`; the mask discards whatever came from past the word. #[inline] pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { let len = word.len(); @@ -53,10 +46,6 @@ pub fn key_and_hash_readable(word: &[u8], readable: usize) -> (u64, u64) { } // SAFETY: `readable >= 8` bytes exist from `word.as_ptr()`, and `len <= 7 < 8`. let raw = unsafe { word.as_ptr().cast::().read_unaligned() }; - // Two table loads, both independent of `raw`, so they overlap its load instead of queueing - // behind a shift chain. Computing them instead -- `u64::MAX >> (64 - 8 * len)` and `len << 56` -- - // measured 1.020x against this 1.034x: the arithmetic is more instructions on a path that is - // bound by how many it runs, and the loads were never waiting on anything. `len <= 7` is guaranteed above, so neither index is checked at runtime. // SAFETY: `len <= INLINE_KEY_BYTES == 7`, and both tables have 8 entries. let (mask, tag) = unsafe { (*KEY_MASK.get_unchecked(len), *LEN_TAG.get_unchecked(len)) }; let key = (raw & mask) | tag; @@ -87,7 +76,7 @@ pub fn key_and_hash(word: &[u8]) -> (u64, u64) { (key, mix(key)) } -/// One probe entry: a digest to confirm the slot, and the id to return. **8 bytes.** +/// #[derive(Clone, Copy, Debug, Default)] #[repr(C)] struct Entry { @@ -98,7 +87,6 @@ struct Entry { const _: () = assert!(size_of::() == 8); -/// `KEY_MASK[len]` keeps the low `len` bytes; `LEN_TAG[len]` is the length in the top byte. static KEY_MASK: [u64; 8] = [ 0x0000_0000_0000_0000, 0x0000_0000_0000_00FF, @@ -120,14 +108,11 @@ static LEN_TAG: [u64; 8] = [ 7 << 56, ]; -/// The 32 bits an entry stores to reject an out-of-vocabulary query. Derived from the key by a #[inline(always)] -/// The 32 bits that confirm a slot really holds the queried token. fn digest_of(hash: u64) -> u32 { (hash >> 32) as u32 } -/// `slot -> (offset into `bytes`, length)`. Off the probe path on purpose: only the reverse lookup #[derive(Clone, Copy, Debug, Default)] struct Span { start: u32, @@ -157,13 +142,13 @@ pub struct BucketVocabStore { mphf: Mphf, /// All token bytes, concatenated. Ordered by MPHF slot. bytes: Box<[u8]>, - /// `entries[slot]` -> (key, id). Ordered by MPHF slot. The probe touches only this. entries: Box<[Entry]>, - /// `spans[slot]` -> where the token's bytes are. Reverse lookup only. spans: Box<[Span]>, /// `id_to_slot[token_id] -> entry_idx` -> index into entries as the entries are not really sorted. id_to_slot: Box<[u32]>, /// Number of real tokens. Cached at build so `len()` is O(1): `entries` is sized to the + /// MPHF's non-minimal slot range (with phantom padding slots), so its length is not the + /// token count. n: usize, } @@ -181,6 +166,7 @@ impl PartialEq for BucketVocabStore { if self.len() != other.len() { return false; } + // early exit as soon as there is a missmatch for id in 0..self.len() { if self.id_to_token(id as u32) != other.id_to_token(id as u32) { return false; @@ -205,6 +191,9 @@ impl BucketVocabStore { .map(|(s, _)| key_and_hash(s.as_slice()).1) .collect(); + // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely + // (~n^2/2^65); if one ever fires, switch the key type to u128. The byte check below makes + // a collision a correct miss at query time, but it would drop a token at build, so guard. let mut seen = HashSet::with_capacity(n); for k in &keys { @@ -220,11 +209,17 @@ impl BucketVocabStore { } } + // 3. Build the (non-minimal) `FastPtrHash` via `PtrHashParams::default_fast()`; query with `.index()`. let params = PtrHashParams::default_fast(); let mphf = Mphf::new(&seen.into_iter().collect::>(), params); + // FastPtrHash is non-minimal: `index()` may return a slot up to `max_index()` (>= n), + // so `entries` must be sized to cover the whole slot range. Slots never written by the + // build loop stay as the default `Entry { len: 0, .. }` (phantom/padding slots), which + // enumeration/count paths filter out via `len > 0`. let n_slots = mphf.max_index(); + // 4. Place each token at its MPHF slot; build the slab and the id->slot reverse table. let total: usize = tokens.iter().map(|(s, _)| s.len()).sum(); let max_id = tokens.iter().map(|(_, id)| *id).max().unwrap(); let mut bytes = Vec::with_capacity(total); @@ -262,6 +257,7 @@ impl BucketVocabStore { } pub fn new() -> Self { + // convenient to build empty edit later. let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), @@ -274,6 +270,9 @@ impl BucketVocabStore { } /// This function is the equivalent of `get` on a HashaMap, it return the id + /// corresponding to the key `q`. Since `mphf` always return a slot, we check + /// whether the token indexed by that slot actually match the query. We don't + /// care about collisions on query because of this! #[inline] pub fn get_bytes(&self, q: &[u8]) -> Option { if self.entries.is_empty() { @@ -286,33 +285,30 @@ impl BucketVocabStore { } /// The id for `q`, together with whether that entry may be folded. One probe and one entry + /// load: the flag is a bit of the id the probe already read. #[inline] pub fn get_bytes_foldable(&self, q: &[u8]) -> Option<(u32, bool)> { let (key, hash) = key_and_hash(q); self.get_keyed_foldable(key, hash) } - /// The slot a hash lands in. Split out of the probe so a caller with many words can issue all #[inline(always)] pub fn probe_slot(&self, hash: u64) -> usize { self.mphf.index(&hash) } - /// The `(key, id)` at a slot, without deciding anything. #[inline(always)] pub fn entry_at(&self, slot: usize) -> (u32, u32) { let e = self.entries[slot]; (e.digest, e.id) } - /// Decide a probe from what [`Self::entry_at`] already loaded. #[inline(always)] pub fn resolve_foldable(hash: u64, entry: (u32, u32)) -> Option<(u32, bool)> { let (edigest, eid) = entry; (edigest == digest_of(hash)).then_some((eid & VOCAB_ID_MASK, eid & FOLD_BIT != 0)) } - /// [`Self::get_bytes_foldable`] for a caller that already has the word's key and hash. #[inline] pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { let _ = key; @@ -352,6 +348,7 @@ impl BucketVocabStore { #[inline] pub fn id_to_token(&self, id: u32) -> Option { + // we are not sure its a valid utf8 so if not, adds replacement char self.id_to_token_bytes(id) .map(|b| String::from_utf8_lossy(b).into_owned()) } @@ -361,6 +358,10 @@ impl BucketVocabStore { } /// One past the highest id this vocabulary can hold. + /// + /// Ids are not dense: a config may leave gaps, so [`Self::len`] counts entries and is *not* an + /// id bound. Anything that walks ids has to bound itself by this and skip the holes, which + /// [`Self::id_to_token_bytes`] reports as `None`. pub fn id_space(&self) -> usize { self.id_to_slot.len() } @@ -374,6 +375,7 @@ impl BucketVocabStore { .iter() .zip(self.spans.iter()) .filter(|(_, sp)| sp.len > 0) + // Mask: the stored id carries FOLD_BIT, which must never escape this type. .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() @@ -390,6 +392,7 @@ impl BucketVocabStore { .iter() .zip(self.spans.iter()) .filter(|(_, sp)| sp.len > 0) + // Mask: the stored id carries FOLD_BIT, which must never escape this type. .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() @@ -463,6 +466,7 @@ mod tests { #[test] fn eq_matches_on_dense_content() { + // Models use dense ids (0..n); equality must reflect the token set on that range. let a = BucketVocabStore::build(vec![(b"x".to_vec(), 0), (b"y".to_vec(), 1)]); let b = BucketVocabStore::build(vec![(b"y".to_vec(), 1), (b"x".to_vec(), 0)]); let c = BucketVocabStore::build(vec![(b"x".to_vec(), 0), (b"z".to_vec(), 1)]); From dd84d58de8fa7ad69b693ea26c50d0641eddd924 Mon Sep 17 00:00:00 2001 From: Lysandre Debut Date: Fri, 7 Aug 2026 08:40:01 +0200 Subject: [PATCH 15/16] wide 1.6.0 is yanked -> update cargo.toml to reference 1.5.0 --- tokenizers/tk-encode/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index bdea45df3..0b7ac82b8 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -60,7 +60,7 @@ memchr = "2.8.2" unicode-normalization = "0.1.25" yada = "0.7.0" libc = "0.2" -wide = "1.6.0" +wide = "1.5.0" # Latest released tokenizers, used as the comparison baseline by the CI benchmark # (examples gated on `bench-baseline`). Optional so production builds never pull it. From b67245a6c326a8a3dfb38a1aab1ab283558e768a Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 18:00:47 +0900 Subject: [PATCH 16/16] perf(bpe): probe the cache before the fold `tokenize_spans` ran `fold_id_keyed` -- an MPHF probe, a pilot load plus a dependent entry load into the whole vocabulary -- ahead of `probe_emit_keyed`, which is one load of the home slot and an unconditional store of its lanes. The expensive probe went first and answered only the words that are their own vocabulary entry, while every word the cache was about to serve paid it for nothing. On a warm cache that is nearly all of them. The cache now goes first and the fold answers the miss, where it still beats running the merge engine. A folded word is inserted, so its second and later occurrences come off the cache instead of re-probing the vocabulary. The order follows whichever probe is cheaper, and here the fused emit already made that the cache. On the branch behind #2313, where `900b6a48`'s digest store makes the fold cheap and the cache is still reached through `lookup_keyed`, the same reordering measures 0.951 -- so it is the relative cost that decides, not the order itself. ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against this branch's head with the LLC evicted between binaries. MB/s before -> after: gpt2 english 1128 -> 1182 code 609 -> 611 dense 1534 -> 1621 chinese 882 -> 902 hindi 544 -> 595 thai 617 -> 654 korean 584 -> 610 russian 708 -> 767 greek 670 -> 710 arabic 621 -> 676 llama-3 english 1090 -> 1141 code 714 -> 730 dense 1412 -> 1490 chinese 914 -> 918 hindi 835 -> 810 thai 932 -> 1003 korean 798 -> 866 russian 898 -> 968 greek 870 -> 936 arabic 904 -> 984 warm geomean 1.053, cold 1.039. Against c7ae7f4 on the same box this takes the branch from 0.925 to 0.982. Reserving two ids per span instead of one, which c7ae7f4 does, measures +0.24% here -- inside the +-0.8% geomean noise floor -- so `975ed8df`'s one-per-span reservation stays. Byte-exact: token counts unchanged on all 20 model x corpus pairs and equal to c7ae7f4's. `prove_fold` only sets the bit for an entry that merging its own text reproduces, so a folded word and a merged word give the same ids; the reorder moves which path answers, not what it answers. 370 tests pass. --- tokenizers/tk-encode/src/models/bpe/model.rs | 50 ++++++++++++++------ 1 file changed, 35 insertions(+), 15 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 08d283f31..6d8395003 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -308,10 +308,6 @@ impl pipeline::Model for PipelineBPE { let bytes = sequence.as_bytes(); let (key, hash) = key_and_hash(bytes); - if let Some(id) = self.fold_id_keyed(key, hash) { - output.push(PipelineToken { id }); - return Ok(()); - } let BpeScratch { symbols, @@ -319,7 +315,8 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // A word seen before costs a probe instead of a merge. + // Cache before fold, for the reason given in `tokenize_spans`: the cache is one load and + // the fold is an MPHF probe, so the fold must not run ahead of it. let insert_at = if let Some(cache) = word_cache.as_mut() { match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { @@ -332,6 +329,16 @@ impl pipeline::Model for PipelineBPE { None }; + if let Some(id) = self.fold_id_keyed(key, hash) { + output.push(PipelineToken { id }); + if let Some(cache) = word_cache.as_mut() + && let Some(at) = insert_at + { + cache.insert(at, std::iter::once(id)); + } + return Ok(()); + } + let start = output.len(); self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -385,18 +392,17 @@ impl pipeline::Model for PipelineBPE { capacity = output.capacity(); } - // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a - // word that is itself a foldable vocabulary entry in one probe, and those words never - // reach the cache. Probing the cache first would populate it with words the fold - // already serves for free, and the two paths would disagree about what it holds. + // The cache goes first. It is one direct-mapped load; the fold is an MPHF probe, which + // is a pilot load plus a dependent entry load into the whole vocabulary. Running the + // fold ahead of the cache paid that on every pre-token including the ones the cache + // was about to answer, and on a warm cache that is nearly all of them. + // + // The two still agree on ids: `prove_fold` only sets the bit for an entry that merging + // its own text reproduces, so a folded word and a merged word give the same answer. + // What changes is that a foldable word now gets *inserted*, so its second and later + // occurrences come off the cache instead of re-probing the vocabulary. let (key, hash) = key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); - if let Some(id) = self.fold_id_keyed(key, hash) { - // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. - unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; - cursor += 1; - continue; - } let mut placement = None; if let Some(cache) = word_cache.as_mut() { @@ -425,6 +431,20 @@ impl pipeline::Model for PipelineBPE { } } + // Cache miss. The fold still answers a word that is its own vocabulary entry in one + // probe, which beats running the merge engine for it. + if let Some(id) = self.fold_id_keyed(key, hash) { + // SAFETY: the check above leaves at least `MAX_INLINE_IDS >= 1` slots past `cursor`. + unsafe { output.as_mut_ptr().add(cursor).write(PipelineToken { id }) }; + cursor += 1; + if let Some(cache) = word_cache.as_mut() + && let Some(at) = placement + { + cache.insert(at, std::iter::once(id)); + } + continue; + } + // SAFETY: `cursor` counts what the fast paths wrote; the merge below uses `output` unsafe { output.set_len(cursor) }; let start = output.len();