From d0b775082258edc11ff8095f24c83d7b9b0880eb Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 12:11:00 +0900 Subject: [PATCH 1/9] 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 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 900b6a480a266748098418166b30d14e0b5c77a1 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 15:32:52 +0900 Subject: [PATCH 6/9] perf(bpe): u64 keys and a digest instead of storing the length The probe already stopped reading the byte slab for short words; this drops what remained. `Entry` becomes `{ digest: u32, id: u32 }` -- 8 bytes, so twice as many per cache line -- and a slot is verified by comparing 32 bits of the hash rather than the key, so nothing stores or compares a length. `LookupKey` follows from `u128` to `u64`. Together with what this branch already had -- one hash serving both the fold and the cache probe, and the key living in the entry so the probe is two dependent loads instead of three -- the fold path is now: one masked load and a multiply for the key, one pilot load, one entry load, one compare. Measured +6.1% geomean over tokbench's 29 gpt2 cells (median of four interleaved runs against gigatoken inside each run). Biggest on the cells with the most pretokens per byte: chat-deepseek 1.25x, added-special-dense 1.16x, chat-chatml 1.16x, agentic-traces 1.16x, chat-llama3 1.16x. Verification becomes probabilistic. A wrong id needs a 32-bit digest collision on a slot that is occupied: (50257/65536) x 2^-32 = 1.8e-10 per DISTINCT pretoken -- per distinct, not per query, since a word keys to the same slot and digest every time. Distinct pretokens grow ~1e4 per MB by Heaps' law, so a 1 TB corpus reaches ~1e7-1e8 and expects 0.018 wrong ids; expecting one needs 5.5e9 distinct unseen byte strings, essentially the whole digest space. Exhaustively checking every 1-, 2- and 3-byte string (16,777,216 arbitrary-byte queries) gives 0 false positives with all 50,257 real tokens resolving. --- tokenizers/tk-encode/src/models/bpe/model.rs | 12 +- tokenizers/tk-encode/src/utils/word_cache.rs | 360 +++++------------ .../tk-encode/src/vocab/bucket_vocab_store.rs | 382 ++++++------------ 3 files changed, 237 insertions(+), 517 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 5882dc0ae..4f60a944e 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -246,10 +246,10 @@ 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)?; + let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } @@ -319,7 +319,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(()); } @@ -332,7 +332,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(()); @@ -395,14 +395,14 @@ 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) { output.push(PipelineToken { id }); continue; } let mut placement = None; 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 })); continue; diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index 335073818..7e857bccd 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -1,70 +1,16 @@ -//! 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; -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; + + +/// 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]`) pub struct WordCache { @@ -87,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; @@ -110,37 +52,81 @@ 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`] 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(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"); + 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. + /// # 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]); + } + } + 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`] so + #[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); 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, @@ -152,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, @@ -161,19 +146,15 @@ 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; - 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 { - // 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(); @@ -202,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 = @@ -248,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 { @@ -268,7 +246,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) { @@ -305,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>), @@ -366,115 +343,35 @@ 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. -/// -/// 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, home slot and tag of a word. #[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), - // ^ 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) + index: (hash & placement_mask) as usize, + tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), } } 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 { + 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), - ); - } - debug.finish() + write!(f, "LookupKey(hash {:#018x})", self.0) + } } } @@ -489,6 +386,15 @@ 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), +} + struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, @@ -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,28 +551,26 @@ 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; assert_ne!(key(b"aaaaaaaaaaaaaa\x7f"), key(b"aaaaaaaaaaaaaa\xff")); assert_ne!(key(b"abcd"), key(b"abcd\0")); - assert_eq!(key(b"aaaaaaaaaaaaaa\xff").0 & LookupKey::TAG_MASK, 0); + 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 - /// 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, @@ -684,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, @@ -693,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, @@ -702,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, @@ -732,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 { @@ -744,26 +633,22 @@ 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. #[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}" ); } } /// 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); @@ -785,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); @@ -795,7 +677,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]); @@ -803,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() { @@ -816,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:?}"); @@ -824,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); @@ -836,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); @@ -852,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); @@ -871,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); @@ -890,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); @@ -902,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); @@ -916,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); @@ -953,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); @@ -967,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 08e077874..188da525b 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,85 +4,66 @@ 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. +static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, 0xA409_3822_299F_31D0, 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]); +type Mphf = FastPtrHash; -/// 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 { - // `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 -}; -/// 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". +/// Bit 31 of a stored id: the token provably encodes to itself, so a pretoken equal to it can be +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 `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. +/// 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 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 +71,6 @@ 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`. 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 +83,51 @@ 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`. -/// -/// 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 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. +/// 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 { - /// 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. +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 +#[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, @@ -163,16 +157,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. + /// `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]>, /// 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, } @@ -190,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; @@ -210,15 +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 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 - // (~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 { @@ -234,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); @@ -256,13 +236,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, @@ -282,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()), @@ -295,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() { @@ -306,64 +282,46 @@ 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) + (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); - 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 + #[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 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. + /// [`Self::get_bytes_foldable`] for a caller that already has the word's key and hash. #[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,14 +345,13 @@ 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] 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()) } @@ -404,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() } @@ -417,12 +370,11 @@ impl BucketVocabStore { } pub fn content(&self) -> Vec<(String, u32)> { - self.spans + self.entries .iter() - .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) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -434,12 +386,11 @@ 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) - // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|(_, m)| m.id & VOCAB_ID_MASK) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) + .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 +400,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)]); @@ -598,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 44f3b4fe63f74e1b4054a86fd3bfc9dd6fda1695 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 17:45:25 +0900 Subject: [PATCH 7/9] perf(bpe): probe the cache before the fold, through the fused emit `tokenize_spans` ran `fold_id_keyed` -- an MPHF probe -- ahead of the word cache, and reached the cache through `lookup_keyed`, which walks the tag window, builds a `&[u32]` and `extend`s it. So the expensive probe ran first and answered only the words that are their own vocabulary entry, while every cache-servable word paid it for nothing. On a warm cache that is nearly all of them. Two changes, and they only pay together: - the cache is reached through `probe_emit_keyed`, already in `word_cache`: a hit is one load of the home slot and an unconditional store of its lanes, written straight at a running cursor into the caller's buffer, so the ids never become a slice and the line is never read twice. `output` is reserved once at two ids per span. - the fold moves *behind* that probe, and a folded word is now inserted, so its second and later occurrences come off the cache instead of re-probing the vocabulary. The order is the whole point, and it follows whichever probe is cheaper. Measured on this branch, against its own head: fused emit, fold still first 1.013 fused emit, cache first 1.110 <- the reorder is +9.2% of that on its own Reordering *without* the fused emit measures 0.951 -- a 5% regression -- because `lookup_keyed` costs more than the fold probe the digest store made cheap. That is why the two land together. ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against the branch head with the LLC evicted between binaries. MB/s before -> after: gpt2 english 786 -> 836 code 392 -> 446 dense 1572 -> 1651 chinese 900 -> 925 hindi 369 -> 507 thai 497 -> 610 korean 570 -> 592 russian 678 -> 702 greek 630 -> 658 arabic 604 -> 603 llama-3 english 794 -> 856 code 418 -> 432 dense 1442 -> 1621 chinese 957 -> 1001 hindi 652 -> 675 thai 803 -> 958 korean 667 -> 823 russian 790 -> 896 greek 776 -> 888 arabic 698 -> 886 warm geomean 1.110 (gpt2 1.095, llama-3 1.126), cold 1.058. Worst cell 0.999, best 1.374. Byte-exact: token counts are unchanged on all 20 model x corpus pairs, and equal to c7ae7f4's on the same corpora. `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 | 75 ++++++++++++++++---- 1 file changed, 60 insertions(+), 15 deletions(-) diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 4f60a944e..7936e3d01 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; @@ -377,9 +377,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% are at most two, so reserve for two apiece + // and emit a cache hit by writing straight at a running cursor: `extend` would re-check + // capacity and re-read the length for every word. + 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 @@ -389,28 +392,66 @@ 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. + // The probe needs somewhere to put the ids before it knows how many there are, so + // make the room first: after this, `MAX_INLINE_IDS` writes past the cursor are + // always inside the allocation. + 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(); + } + + // Cache first, and through the fused probe: a hit is one load of the home slot and an + // unconditional store of its lanes, written straight at the cursor, so the ids never + // become a slice and the line is never read twice. That makes the cache cheaper than + // the fold's MPHF probe, which is what lets the fold move behind it. + // + // 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. 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 }); - continue; - } let mut placement = None; if let Some(cache) = word_cache.as_mut() { - match cache.lookup_keyed(key, hash) { - Lookup::Hit(ids) => { + // SAFETY: the check above leaves `MAX_INLINE_IDS` slots past `cursor`. + let found = unsafe { + cache.probe_emit_keyed(key, hash, output.as_mut_ptr().add(cursor).cast::()) + }; + match found { + ProbeEmit::Wrote(n) => { + cursor += n; + continue; + } + 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), + } + } + + // 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` + // 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 @@ -422,7 +463,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 above. + unsafe { output.set_len(cursor) }; Ok(()) } From 2ab49f9eb6aebe4ae171fb20426454aafea344c7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 18:43:47 +0900 Subject: [PATCH 8/9] perf(added-vocab): search the whole prefix, not just its first byte The single-bucket scan looked for `prefix[0]` and rejected each candidate by hand, restarting `memchr` on `&bytes[pos + 1..]` every time. One byte of a long needle is a poor filter on text that is dense in that byte and holds no match: `<|endoftext|>` over a corpus full of `<|xs0|>` stops at every `<` and dies at the third byte, and every restart pays memchr's prologue again. Measured over the same one-bucket vocabulary in one process, so the two arms share a build: on a 9.8%-`<` corpus that matches nothing, scanning the first byte cost 0.58 ns/B against 0.12 for `memmem`, which picks a *rare* byte of the needle instead -- 4.3x. Where candidates are already sparse it costs at most 0.015 ns/B (english 0.009 -> 0.024), so it trades a little on inputs where the scan is already free for a lot on the inputs where it is not. `nibble_mask_match` already avoids this for two or more buckets, where it is called the restart penalty; the one-bucket path never got it. `find_iter` yields the same positions in the same order, so the leftmost match is unchanged and `match_fast` still confirms the length sub-list. --- tokenizers/tk-encode/src/vocab/buckets.rs | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/tokenizers/tk-encode/src/vocab/buckets.rs b/tokenizers/tk-encode/src/vocab/buckets.rs index 9a78a8cec..ff46cf50c 100644 --- a/tokenizers/tk-encode/src/vocab/buckets.rs +++ b/tokenizers/tk-encode/src/vocab/buckets.rs @@ -333,14 +333,23 @@ impl Buckets { // needle = the bucket's shared first byte. Assumes a non-empty prefix // (false only if a lone 1-byte token is the sole holder of its first byte); store // the first byte explicitly if that case ever appears. - let needle = self.buckets[0].prefix[0]; - let mut search = 0usize; - while let Some(off) = memchr::memchr(needle, &bytes[search..]) { - let pos = search + off; + // Search for the bucket's whole shared prefix, not just its first byte. + // + // One byte of a long needle is a poor filter on text that is dense in that byte and + // holds no match: `<|endoftext|>` over a corpus full of `<|xs0|>` stops at every `<` + // and dies at the third byte. Restarting `memchr` on `&bytes[pos + 1..]` per + // candidate then pays its prologue again each time. Measured over the same + // one-bucket vocabulary in one process, that cost 0.58 ns/B on a 9.8%-`<` corpus + // that matches nothing, against 0.12 for `memmem`, which picks a *rare* byte of the + // needle instead. Where candidates are already sparse it costs at most 0.015 ns/B. + // + // `nibble_mask_match` already avoids this for two or more buckets and calls it the + // restart penalty. Same positions in the same order, so the leftmost match is + // unchanged and `match_fast` still confirms the length sub-list. + for pos in memchr::memmem::find_iter(bytes, &self.buckets[0].prefix) { if let Some((id, len)) = self.match_fast(bytes, pos, 0) { return Some((id, pos as u32, len)); } - search = pos + 1; } None } From 3a2787baf7fe14207ee312386fe7de9b0006c090 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 21:01:46 +0900 Subject: [PATCH 9/9] fix(vocab): confirm a probe on the whole hash, not 32 bits of it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Entry` stored `digest_of(hash) = (hash >> 32) as u32` and nothing else checked the hit -- there is no byte comparison against the token. The MPHF slot is derived from that same `hash`, so the slot and the digest are correlated rather than independent, which makes the real false-positive rate for an out-of-vocabulary query worse than the 2^-32 the width suggests. When it fires, the vocabulary answers a pretoken that is not in it with some other token's id, silently. It is not hypothetical. With a different (weaker) key hash it fired on the first corpus tried: ` dignified`, which gpt2 does not contain, came back as `Ġsignifies` (id 43854) instead of `[13469, 1431]` -- both ten bytes long. So the design is one hash-function change away from mis-tokenizing real text, and the failure is invisible without an oracle to compare against. Store the whole 64 bits and compare all of them: a false positive becomes 2^-64, the standard the word cache already holds long words to. `digest_of` goes with it. `Entry` grows 8 -> 16 bytes, on a path the profile puts at roughly a quarter of encode time, so the cost was worth measuring rather than assuming: over 29 tokbench gpt2 cells, 3 interleaved rounds, medians, it is **1.0066x (12/29 faster)** -- free, within noise. All 30 corpora stay byte-exact against the reference tokenizer. --- .../tk-encode/src/vocab/bucket_vocab_store.rs | 38 ++++++++++--------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 188da525b..cd9a8ee4f 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -87,16 +87,25 @@ 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.** +/// One probe entry: the whole hash to confirm the slot, and the id to return. **16 bytes.** +/// +/// A 32-bit digest cut from `hash >> 32` was not enough. The MPHF slot comes from the same `hash`, +/// so slot and digest are correlated rather than independent, and nothing else confirms the hit -- +/// there is no byte comparison against the token. A pretoken that is *not* in the vocabulary can +/// therefore land on a slot whose digest also matches and be answered with another token's id. +/// Not hypothetical: with a different (weaker) key hash it fired on the first corpus tried, emitting +/// `Ġsignifies` for the pretoken ` dignified`, both ten bytes long. This is a silent-wrong-answer +/// path, so it is worth eight bytes an entry to make a false positive 2^-64 -- the standard the word +/// cache already holds long words to. #[derive(Clone, Copy, Debug, Default)] #[repr(C)] struct Entry { - digest: u32, + hash: u64, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } -const _: () = assert!(size_of::() == 8); +const _: () = assert!(size_of::() == 16); /// `KEY_MASK[len]` keeps the low `len` bytes; `LEN_TAG[len]` is the length in the top byte. static KEY_MASK: [u64; 8] = [ @@ -120,13 +129,6 @@ 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 { @@ -240,7 +242,7 @@ impl BucketVocabStore { let (key, hash) = key_and_hash(s.as_slice()); let slot = mphf.index(&hash); entries[slot] = Entry { - digest: digest_of(hash), + hash, id: *id, }; spans[slot] = Span { @@ -282,7 +284,7 @@ impl BucketVocabStore { let (key, hash) = key_and_hash(q); let slot = self.mphf.index(&hash); let e = self.entries[slot]; - (e.digest == digest_of(hash)).then_some(e.id & VOCAB_ID_MASK) + (e.hash == 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 @@ -300,16 +302,16 @@ impl BucketVocabStore { /// The `(key, id)` at a slot, without deciding anything. #[inline(always)] - pub fn entry_at(&self, slot: usize) -> (u32, u32) { + pub fn entry_at(&self, slot: usize) -> (u64, u32) { let e = self.entries[slot]; - (e.digest, e.id) + (e.hash, 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)) + pub fn resolve_foldable(hash: u64, entry: (u64, u32)) -> Option<(u32, bool)> { + let (ehash, eid) = entry; + (ehash == 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. @@ -321,7 +323,7 @@ impl BucketVocabStore { } let slot = self.mphf.index(&hash); let e = self.entries[slot]; - (e.digest == digest_of(hash)).then_some((e.id & VOCAB_ID_MASK, e.id & FOLD_BIT != 0)) + (e.hash == 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.