perf(bpe): Last missing nits to get up to speed

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};
 
 const GATE_MULTI: u16 = 8;
 const GATE_ASCII: u16 = 24;
@@ -240,13 +240,16 @@ 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.
+    /// 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(&self, sequence: &str) -> Option<u32> {
+    fn fold_id_keyed(&self, key: u64, hash: u64) -> Option<u32> {
         // 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,7 +312,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 (key, hash) = key_and_hash(bytes);
+        if let Some(id) = self.fold_id_keyed(key, hash) {
             output.push(PipelineToken { id });
             return Ok(());
         }
@@ -322,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(sequence.as_bytes()) {
+            match cache.lookup_keyed(key, hash) {
                 Lookup::Hit(ids) => {
                     output.extend(ids.iter().map(|&id| PipelineToken { id }));
                     return Ok(());
@@ -367,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
@@ -379,26 +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.
-            if let Some(id) = self.fold_id(sequence) {
-                output.push(PipelineToken { id });
-                continue;
+            // 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);
+
             let mut placement = None;
             if let Some(cache) = word_cache.as_mut() {
-                match cache.lookup(sequence.as_bytes()) {
-                    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::<u32>())
+                };
+                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
@@ -410,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(())
     }
 

tokenizers/tk-encode/src/tokenizer/pipeline.rs

@@ -994,6 +994,24 @@ impl PipelineTokenizer {
         add_special_tokens: bool,
     ) -> Result<Vec<PipelineToken>> {
         let mut output = Vec::with_capacity(input.len() / 4);
+        self.encode_generic_into::<STAGE>(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<const STAGE: u8>(
+        &self,
+        input: &str,
+        add_special_tokens: bool,
+        output: &mut Vec<PipelineToken>,
+    ) -> 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(())
     }
 }