From 25c5acb95d2629d4884ed158cbe6b29698a1aca2 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 12:42:28 +0900 Subject: [PATCH] perf(bpe): put the key in the entry, so the fold probe is two loads 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 675cfc18f..2bc0dd74c 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}; 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)]);