From 2296c0b20e08cc944e1620569e84ef22ae3c4829 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 7 Aug 2026 12:35:35 +0900 Subject: [PATCH] perf(bpe): pack a short word into its key instead of hashing it 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