perf(tk-encode): SIMD-first encode pipeline — fused split cache + de-virtualized BPE

tokenizers/atomsplit/src/fsm.rs

@@ -502,8 +502,38 @@ pub fn fsm_deepseek(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
 /// `\p{N}{1,3}` cap (numbers are unbounded) and no `\s*[\r\n]`/trailing-`[\r\n]*` rules.
 /// ┌── OWNER: shared (scalar) ──┐
 #[must_use]
+/// Pack a span's first ≤15 bytes into a `u128` (length in the top byte) via one unaligned 16-byte
+/// load — the bytes the FSM just walked, still hot. `0` = not packable (empty or >15 bytes), so a
+/// caller keyed cache must byte-verify those (two long spans sharing a 15-byte prefix would alias).
+#[inline(always)]
+pub fn pack_key(text: &[u8], start: usize, len: usize) -> u128 {
+    if len == 0 || len > 15 {
+        return 0;
+    }
+    let v: u128 = if start + 16 <= text.len() {
+        // SAFETY: `start + 16 <= text.len()` — 16 readable bytes; unaligned load is always valid.
+        unsafe { (text.as_ptr().add(start) as *const u128).read_unaligned() }
+    } else {
+        let mut b = [0u8; 16];
+        b[..len].copy_from_slice(&text[start..start + len]);
+        u128::from_le_bytes(b)
+    };
+    (v & ((1u128 << (8 * len)) - 1)) | ((len as u128) << 120)
+}
+
 pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
+    byte_level_impl::<false>(text, tags, out, &mut [])
+}
+
+#[inline(always)]
+fn byte_level_impl<const KEYED: bool>(
+    text: &[u8],
+    tags: &[u8],
+    out: &mut [Span],
+    keys: &mut [u128],
+) -> usize {
     debug_assert!(out.len() >= text.len() && tags.len() >= text.len());
+    debug_assert!(!KEYED || keys.len() >= text.len());
     const LET: u8 = Atom::Letter as u8;
     const NW: u8 = Atom::NumWord as u8;
     const NO: u8 = Atom::NumOther as u8;
@@ -576,6 +606,9 @@ pub fn fsm_byte_level(text: &[u8], tags: &[u8], out: &mut [Span]) -> usize {
             _ => i += char_len(text[i]),
         }
         out[w] = (start as u32, i as u32);
+        if KEYED {
+            keys[w] = pack_key(text, start, i - start); // packed right at the bound, bytes hot
+        }
         w += 1;
     }
     w

tokenizers/tk-encode/src/models/bpe/flat_cache.rs

@@ -34,17 +34,22 @@ fn clear_on_full() -> bool {
     *C.get_or_init(|| std::env::var("CACHE_EVICT").map(|v| v == "clear").unwrap_or(false))
 }
 
-/// 24 bytes: hash(8) + koff(4) + ioff(4) + klen(2) + ilen(2) + freq(1) + pad.
+/// hash(8) + key(16) + koff(4) + ioff(4) + klen(2) + ilen(2) + freq(1). `key != 0` = a ≤15-byte
+/// pre-token packed inline (length in top byte): a hit is a register 128-bit compare, no `kbytes`
+/// load, no `memcmp`. `key == 0` = a long (>15B) pre-token, verified via `kbytes`. The ids always
+/// live in the `ids` arena (inline-in-slot was measured slower in our L2-resident regime — the
+/// wider slot table costs more than the saved arena load, which is already hot here).
 #[derive(Clone, Copy)]
 struct CSlot {
     hash: u64,
+    key: u128,
     koff: u32,
     ioff: u32,
     klen: u16,
     ilen: u16,
     freq: u8,
 }
-const EMPTY: CSlot = CSlot { hash: 0, koff: 0, ioff: 0, klen: 0, ilen: 0, freq: 0 };
+const EMPTY: CSlot = CSlot { hash: 0, key: 0, koff: 0, ioff: 0, klen: 0, ilen: 0, freq: 0 };
 
 pub(crate) struct FlatCache {
     slots: Box<[CSlot]>,
@@ -137,6 +142,7 @@ impl FlatCache {
             }
             slots2[i] = CSlot {
                 hash: s.hash,
+                key: s.key,
                 koff,
                 ioff,
                 klen: s.klen,
@@ -163,10 +169,14 @@ impl FlatCache {
         }
     }
 
-    /// Lookup + bump the entry's frequency on a hit (so the cull can tell hot
-    /// from one-shot). Byte-verify unless hash-only. Needs `&mut self`.
+    /// Probe for `p`; on a hit bump the entry's frequency (so the cull can tell hot from one-shot)
+    /// and return its ids' `(offset, len)` into the arena; on a miss return `None`. Deliberately
+    /// SMALL so it inlines into the caller's per-pre-token loop — the caller does the memcpy emit
+    /// (`ids_slice` + `extend_from_slice`) so the whole hot path stays branch-local and register-hot.
+    /// `key != 0`: confirm with a register 128-bit compare (`s.key == key`), no `kbytes`, no
+    /// `memcmp`. `key == 0`: byte-verify via `kbytes` (unless hash-only).
     #[inline]
-    pub(crate) fn get(&mut self, p: &[u8], h: u64) -> Option<(u32, u16)> {
+    pub(crate) fn get(&mut self, p: &[u8], h: u64, key: u128) -> Option<(u32, u16)> {
         let mut i = (h as usize) & self.mask;
         loop {
             let s = self.slots[i];
@@ -176,11 +186,16 @@ impl FlatCache {
                 }
                 return None;
             }
-            if s.hash == h
-                && (!self.verify
-                    || (s.klen as usize == p.len()
-                        && self.kbytes[s.koff as usize..s.koff as usize + s.klen as usize] == *p))
-            {
+            let confirmed = if key != 0 {
+                s.key == key // register compare — no arena, no memcmp
+            } else {
+                s.hash == h
+                    && (!self.verify
+                        || (s.klen as usize == p.len()
+                            && self.kbytes[s.koff as usize..s.koff as usize + s.klen as usize]
+                                == *p))
+            };
+            if confirmed {
                 self.slots[i].freq = self.slots[i].freq.saturating_add(1);
                 if self.stats {
                     HITS.fetch_add(1, Ordering::Relaxed);
@@ -191,13 +206,15 @@ impl FlatCache {
         }
     }
 
-    #[inline]
-    pub(crate) fn ids_slice(&self, off: u32, len: u16) -> &[u32] {
-        &self.ids[off as usize..off as usize + len as usize]
+    /// Arena ids for a `(offset, len)` from [`get`], reinterpreted as output tokens for a memcpy emit.
+    #[inline(always)]
+    pub(crate) fn ids_tokens(&self, off: u32, len: u16) -> &[crate::tokenizer::pipeline::PipelineToken] {
+        let (o, n) = (off as usize, len as usize);
+        crate::tokenizer::pipeline::ids_as_tokens(&self.ids[o..o + n])
     }
 
     #[inline]
-    pub(crate) fn insert(&mut self, p: &[u8], h: u64, ids: &[u32]) {
+    pub(crate) fn insert(&mut self, p: &[u8], h: u64, key: u128, ids: &[u32]) {
         if ids.is_empty() || p.len() > u16::MAX as usize || ids.len() > u16::MAX as usize {
             return;
         }
@@ -208,6 +225,7 @@ impl FlatCache {
             self.cull();
         }
         let (koff, ioff) = (self.kbytes.len() as u32, self.ids.len() as u32);
+        // kbytes retained for all (cull compacts by koff/klen); packed-key gets never read it.
         self.kbytes.extend_from_slice(p);
         self.ids.extend_from_slice(ids);
         let mut i = (h as usize) & self.mask;
@@ -216,6 +234,7 @@ impl FlatCache {
         }
         self.slots[i] = CSlot {
             hash: h,
+            key,
             koff,
             ioff,
             klen: p.len() as u16,
@@ -225,3 +244,31 @@ impl FlatCache {
         self.count += 1;
     }
 }
+
+#[cfg(test)]
+mod cache_tests {
+    use super::*;
+
+    #[test]
+    fn slot_fits_one_cache_line() {
+        assert!(std::mem::size_of::<CSlot>() <= 64, "CSlot = {} B", std::mem::size_of::<CSlot>());
+    }
+
+    // Short and long id-runs both round-trip byte-exact through get + ids_tokens.
+    #[test]
+    fn get_roundtrip() {
+        let mut c = FlatCache::new();
+        c.retarget(1);
+        let short: &[u32] = &[10, 20, 30, 40];
+        let long: &[u32] = &[1, 2, 3, 4, 5, 6, 7];
+        let (hs, hl) = (c.hash(b"short"), c.hash(b"longer_token"));
+        c.insert(b"short", hs, 0, short);
+        c.insert(b"longer_token", hl, 0, long);
+        for (p, h, want) in [(&b"short"[..], hs, short), (&b"longer_token"[..], hl, long)] {
+            let (off, len) = c.get(p, h, 0).unwrap_or_else(|| panic!("miss on {p:?}"));
+            let got: Vec<u32> = c.ids_tokens(off, len).iter().map(|t| t.id).collect();
+            assert_eq!(got, want);
+        }
+        assert!(c.get(b"nope", c.hash(b"nope"), 0).is_none()); // absent misses
+    }
+}

tokenizers/tk-encode/src/models/bpe/model.rs

@@ -1271,7 +1271,12 @@ impl PipelineBPE {
 }
 
 impl pipeline::Model for PipelineBPE {
-    fn tokenize_pipeline(&self, sequence: &str, output: &mut Vec<PipelineToken>) -> Result<()> {
+    fn tokenize_pipeline(
+        &self,
+        sequence: &str,
+        output: &mut Vec<PipelineToken>,
+        carried_key: Option<u128>,
+    ) -> Result<()> {
         if sequence.is_empty() {
             return Ok(());
         }
@@ -1291,31 +1296,111 @@ impl pipeline::Model for PipelineBPE {
                 let mut ids = o.borrow_mut();
                 ids.clear();
                 self.encode_piece(sequence, &mut ids);
-                output.extend(ids.iter().map(|&id| PipelineToken { id }));
+                output.extend_from_slice(crate::tokenizer::pipeline::ids_as_tokens(&ids));
             });
             return Ok(());
         }
         PIPE_FLAT_CACHE.with(|cell| {
             let mut cache = cell.borrow_mut();
             cache.retarget(self.cache_id);
-            let h = cache.hash(p);
-            if let Some((off, len)) = cache.get(p, h) {
-                output.extend(cache.ids_slice(off, len).iter().map(|&id| PipelineToken { id }));
+            // Pack the ≤15-byte key from `p` right here — `p` is a slice of the (hot) input, so no
+            // keys buffer to carry and no split penalty. Then a cheap CRC bucket + a register
+            // 128-bit compare in the cache: no ahash, no `kbytes` memcmp. `key == 0` (long/non-GPT)
+            // falls back to hashing the bytes + byte-verify. `carried_key` (if the split fused one)
+            // is preferred to avoid re-packing, but the fallback pack is what removes the keys buffer.
+            let key = carried_key.unwrap_or(0);
+            let h = if key != 0 {
+                crate::tokenizer::pipeline::crc_key_hash(key)
+            } else {
+                cache.hash(p)
+            };
+            if let Some((off, len)) = cache.get(p, h, key) {
+                output.extend_from_slice(cache.ids_tokens(off, len)); // memcpy emit
                 return;
             }
             // Miss: merge into reused scratch, cache the ids, emit.
             PIPE_OUT_IDS.with(|o| {
                 let mut ids = o.borrow_mut();
                 ids.clear();
                 self.encode_piece(sequence, &mut ids);
-                cache.insert(p, h, &ids);
-                output.extend(ids.iter().map(|&id| PipelineToken { id }));
+                cache.insert(p, h, key, &ids);
+                output.extend_from_slice(crate::tokenizer::pipeline::ids_as_tokens(&ids));
             });
         });
         Ok(())
     }
 }
 
+impl PipelineBPE {
+    /// Fused per-chunk tokenize — the de-virtualized hot path. `tokenize_pipeline` was called once
+    /// per pre-token through the `PipelineModel` enum, so every pre-token paid an enum match, a
+    /// thread-local `PIPE_FLAT_CACHE` borrow, and a `retarget` check. Here those happen ONCE per
+    /// chunk; the inner loop then inlines pack_key → hash → probe → emit with no call or TLS barrier
+    /// on the hot (cache-hit) path. Byte-identical to looping `tokenize_pipeline` per pre-token.
+    pub(crate) fn tokenize_chunk(
+        &self,
+        chunk: &str,
+        pre_tokens: &[crate::tokenizer::pipeline::Split],
+        output: &mut Vec<PipelineToken>,
+    ) -> Result<()> {
+        use crate::tokenizer::pipeline::{crc_key_hash, ids_as_tokens, Model as _};
+        let nb = chunk.as_bytes();
+        // Cache off (measurement): defer to the per-pre-token path unchanged.
+        if cache_disabled() {
+            for pt in pre_tokens {
+                let r = pt.range();
+                let key = atomsplit::fsm::pack_key(nb, r.start, r.len());
+                self.tokenize_pipeline(&chunk[r], output, Some(key))?;
+            }
+            return Ok(());
+        }
+        let min_len = min_cache_len();
+        PIPE_FLAT_CACHE.with(|cell| {
+            let mut cache = cell.borrow_mut(); // one borrow for the whole chunk
+            cache.retarget(self.cache_id); // once, not per pre-token
+            for pt in pre_tokens {
+                let r = pt.range();
+                let seq = &chunk[r.start..r.end];
+                if seq.is_empty() {
+                    continue;
+                }
+                let p = seq.as_bytes();
+                if self.ignore_merges {
+                    if let Some(id) = self.vocab.get_bytes(p) {
+                        output.push(PipelineToken { id });
+                        continue;
+                    }
+                }
+                // Short pre-tokens: re-merge is cheaper than a cache round-trip (emit, no insert).
+                if p.len() < min_len {
+                    PIPE_OUT_IDS.with(|o| {
+                        let mut ids = o.borrow_mut();
+                        ids.clear();
+                        self.encode_piece(seq, &mut ids);
+                        output.extend_from_slice(ids_as_tokens(&ids));
+                    });
+                    continue;
+                }
+                let key = atomsplit::fsm::pack_key(nb, r.start, r.len());
+                let h = if key != 0 { crc_key_hash(key) } else { cache.hash(p) };
+                if let Some((off, len)) = cache.get(p, h, key) {
+                    output.extend_from_slice(cache.ids_tokens(off, len)); // memcpy emit
+                    continue;
+                }
+                // Miss: merge into reused scratch, cache the ids, emit.
+                PIPE_OUT_IDS.with(|o| {
+                    let mut ids = o.borrow_mut();
+                    ids.clear();
+                    self.encode_piece(seq, &mut ids);
+                    cache.insert(p, h, key, &ids);
+                    output.extend_from_slice(ids_as_tokens(&ids));
+                });
+            }
+        });
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1848,7 +1933,7 @@ mod tests {
 
         fn pipeline_ids(model: &PipelineBPE, sequence: &str) -> Vec<u32> {
             let mut out = Vec::new();
-            pipeline::Model::tokenize_pipeline(model, sequence, &mut out).unwrap();
+            pipeline::Model::tokenize_pipeline(model, sequence, &mut out, None).unwrap();
             out.iter().map(|t| t.id).collect()
         }
 

tokenizers/tk-encode/src/models/unigram/model.rs

@@ -510,6 +510,7 @@ impl pipeline::Model for Unigram {
         &self,
         sequence: &str,
         output: &mut Vec<pipeline::PipelineToken>,
+        _carried_key: Option<u128>,
     ) -> Result<()> {
         let str_tokens = self.encode(sequence)?;
 

tokenizers/tk-encode/src/models/wordlevel/mod.rs

@@ -212,6 +212,7 @@ impl pipeline::Model for WordLevel {
         &self,
         sequence: &str,
         output: &mut Vec<pipeline::PipelineToken>,
+        _carried_key: Option<u128>,
     ) -> Result<()> {
         if let Some(&id) = self.vocab.get(sequence) {
             output.push(PipelineToken { id })

tokenizers/tk-encode/src/models/wordpiece/mod.rs

@@ -352,6 +352,7 @@ impl pipeline::Model for PipelineWordPiece {
         &self,
         sequence: &str,
         output: &mut Vec<pipeline::PipelineToken>,
+        _carried_key: Option<u128>,
     ) -> Result<()> {
         let mut candidate = String::with_capacity(self.max_input_chars_per_word);
         let mut candidate_tokens = Vec::with_capacity(sequence.len());

tokenizers/tk-encode/src/pre_tokenizers/split.rs

@@ -215,6 +215,7 @@ impl pipeline::PreTokenizer for Split {
         pipeline::split_matches(out, matches, self.behavior);
         Ok(())
     }
+
 }
 
 #[cfg(test)]