perf(bpe): hand the splitter's spans to the model in tiles

tokenizers/bitsplit/src/lib.rs

@@ -34,7 +34,7 @@ mod simd;
 
 pub use models::deepseek::bitsplit_deepseek;
 pub use models::cl100k::{bitsplit_cl100k, bitsplit_qwen};
-pub use models::gpt2::bitsplit_byte_level;
+pub use models::gpt2::{bitsplit_byte_level, bitsplit_byte_level_tiled};
 pub use models::kimi::bitsplit_kimi;
 pub use models::o200k::bitsplit_o200k;
 pub use models::tekken::bitsplit_tekken;
@@ -428,6 +428,81 @@ pub(crate) fn letter_match(tags: &[u8], p: usize, re: usize) -> usize {
 }
 
 
+/// Where the emit loop puts the spans it finds.
+///
+/// Two implementations, and the reason this is a trait at all: the whole-chunk sink writes into one
+/// array sized for the worst case, which for a 1 MB input is ~230k spans = 1.8 MB written and then
+/// read back by the consumer -- 3.5 bytes of span traffic per input byte, none of which stays in
+/// cache. The tiled sink hands each small batch straight to the consumer instead, so the array never
+/// grows past L1 and nothing materialises per document. Both share one emit implementation, because
+/// the emit's control flow (contractions especially) is the last thing worth duplicating.
+pub trait SpanSink {
+    fn push(&mut self, span: Span);
+}
+
+/// The whole-chunk sink: `out` must hold the worst case of one span per byte, plus one.
+pub struct SliceSink<'a> {
+    out: &'a mut [Span],
+    written: usize,
+}
+
+impl<'a> SliceSink<'a> {
+    #[inline(always)]
+    pub fn new(out: &'a mut [Span]) -> Self {
+        Self { out, written: 0 }
+    }
+    #[inline(always)]
+    pub fn written(&self) -> usize {
+        self.written
+    }
+}
+
+impl SpanSink for SliceSink<'_> {
+    #[inline(always)]
+    fn push(&mut self, span: Span) {
+        self.out[self.written] = span;
+        self.written += 1;
+    }
+}
+
+/// A fixed tile handed to `consume` whenever it fills. Flushing from inside the emit's own control
+/// flow is what makes this need no resumable state: a contraction chain (`y'all'd've`) can emit an
+/// unbounded run of spans from a single block, so a caller-driven "emit N then return" API would
+/// have to be resumable *mid-chain*, and getting that wrong changes tokenization.
+pub struct TileSink<'a, F: FnMut(&[Span])> {
+    tile: &'a mut [Span],
+    len: usize,
+    consume: F,
+}
+
+impl<'a, F: FnMut(&[Span])> TileSink<'a, F> {
+    #[inline(always)]
+    pub fn new(tile: &'a mut [Span], consume: F) -> Self {
+        assert!(!tile.is_empty(), "a tile needs room for at least one span");
+        Self { tile, len: 0, consume }
+    }
+
+    /// Hand over whatever is buffered. Must be called once the emit returns, for the tail.
+    #[inline]
+    pub fn flush(&mut self) {
+        if self.len > 0 {
+            (self.consume)(&self.tile[..self.len]);
+            self.len = 0;
+        }
+    }
+}
+
+impl<F: FnMut(&[Span])> SpanSink for TileSink<'_, F> {
+    #[inline(always)]
+    fn push(&mut self, span: Span) {
+        self.tile[self.len] = span;
+        self.len += 1;
+        if self.len == self.tile.len() {
+            self.flush();
+        }
+    }
+}
+
 // ── emit ────────────────────────────────────────────────────────────────────────────────────────
 
 /// `starts` bitmap → spans. Each set bit closes the previous token and opens the next; `tzcnt`
@@ -460,25 +535,24 @@ pub(crate) fn emit(starts: &[u64], nblk: usize, n: usize, out: &mut [Span]) -> u
 ///
 /// A matched contraction overrides the algebra outright: it emits its own span and skips every
 /// start bit inside it, which is how the letter alternative loses the tie (`'sx` → `'s`, `x`).
-pub(crate) fn emit_contr(
+pub(crate) fn emit_contr<S: SpanSink>(
     text: &[u8],
     starts: &[u64],
     flag: &[u64],
     nblk: usize,
     n: usize,
     ci: bool,
-    out: &mut [Span],
-) -> usize {
-    let (mut w, mut open, mut skip) = (0usize, u32::MAX, 0usize);
+    sink: &mut S,
+) {
+    let (mut open, mut skip) = (u32::MAX, 0usize);
     for bi in 0..nblk {
         let mut m = starts[bi];
         let f = flag[bi];
         if f == 0 && skip <= bi * 64 {
             while m != 0 {
                 let pos = (bi * 64 + m.trailing_zeros() as usize) as u32;
                 if open != u32::MAX {
-                    out[w] = Span::new(open, pos);
-                    w += 1;
+                    sink.push(Span::new(open, pos));
                 }
                 open = pos;
                 m &= m - 1;
@@ -493,8 +567,7 @@ pub(crate) fn emit_contr(
                 continue;
             }
             if open != u32::MAX {
-                out[w] = Span::new(open, pos as u32);
-                w += 1;
+                sink.push(Span::new(open, pos as u32));
             }
             open = pos as u32;
             if f >> j & 1 != 0 {
@@ -507,8 +580,7 @@ pub(crate) fn emit_contr(
                     if l == 0 {
                         break;
                     }
-                    out[w] = Span::new(p as u32, (p + l) as u32);
-                    w += 1;
+                    sink.push(Span::new(p as u32, (p + l) as u32));
                     p += l;
                 }
                 if p > pos {
@@ -522,10 +594,8 @@ pub(crate) fn emit_contr(
         }
     }
     if open != u32::MAX {
-        out[w] = Span::new(open, n as u32);
-        w += 1;
+        sink.push(Span::new(open, n as u32));
     }
-    w
 }
 
 

tokenizers/bitsplit/src/models/cl100k.rs

@@ -265,5 +265,7 @@ fn cl100k(
         code = last_code;
         prev_cont = b.cont;
     }
-    emit_contr(text, starts, flag, nblk, ntext, true, out)
+    let mut sink = crate::SliceSink::new(out);
+    emit_contr(text, starts, flag, nblk, ntext, true, &mut sink);
+    sink.written()
 }

tokenizers/bitsplit/src/models/gpt2.rs

@@ -6,7 +6,7 @@
 //! literal alternation that outranks every other arm is miserable in bit algebra and trivial there.
 
 use crate::{
-    AUX_NONE, CODE_CONT, CONT, Span, build_block, emit_contr, to_lead,
+    AUX_NONE, CODE_CONT, CONT, SliceSink, Span, SpanSink, TileSink, build_block, emit_contr, to_lead,
 };
 
 /// Atom tag → dense 3-bit code, shared by both grammars. Unlike deepseek's table, `Mark` is NOT a
@@ -86,14 +86,42 @@ pub fn bitsplit_byte_level(
     flag: &mut [u64],
     out: &mut [Span],
 ) -> usize {
+    let mut sink = SliceSink::new(out);
+    bitsplit_byte_level_sink(text, tags, starts, flag, &mut sink);
+    sink.written()
+}
+
+/// GPT-2 / byte-level pre-tokenization into a [`TileSink`]: the spans go to the consumer in small
+/// batches instead of one per-document array. See [`SpanSink`] for why that is worth a function.
+///
+/// The mask build is unchanged and still one pass over the chunk -- only where the spans *go*
+/// differs, so the grammar (contractions, the `\s+(?!\S)` steal, block-edge carry) is untouched.
+pub fn bitsplit_byte_level_tiled<F: FnMut(&[Span])>(
+    text: &[u8],
+    tags: &[u8],
+    starts: &mut [u64],
+    flag: &mut [u64],
+    tile: &mut [Span],
+    consume: F,
+) {
+    let mut sink = TileSink::new(tile, consume);
+    bitsplit_byte_level_sink(text, tags, starts, flag, &mut sink);
+    sink.flush();
+}
+
+fn bitsplit_byte_level_sink<S: SpanSink>(
+    text: &[u8],
+    tags: &[u8],
+    starts: &mut [u64],
+    flag: &mut [u64],
+    sink: &mut S,
+) {
     let ntext = text.len();
     if ntext == 0 {
-        return 0;
+        return;
     }
     let nblk = ntext.div_ceil(64);
-    assert!(
-        tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk && out.len() >= ntext
-    );
+    assert!(tags.len() >= ntext && starts.len() >= nblk && flag.len() >= nblk);
     let (mut code, mut prev_cont) = (CODE_CONT, 0u64);
 
     for bi in 0..nblk {
@@ -157,6 +185,6 @@ pub fn bitsplit_byte_level(
         code = last_code;
         prev_cont = b.cont;
     }
-    emit_contr(text, starts, flag, nblk, ntext, false, out)
+    emit_contr(text, starts, flag, nblk, ntext, false, sink);
 }
 

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

@@ -373,7 +373,18 @@ impl pipeline::Model for PipelineBPE {
             word_cache,
         } = scratch;
 
-        output.reserve(spans.len() + MAX_INLINE_IDS);
+        // Reserve what the *fast* path could write, so the emit loop's capacity test almost never
+        // fires: a cache hit writes at most `MAX_INLINE_IDS` lanes, so `spans.len() *
+        // MAX_INLINE_IDS` covers every one of them. Only for a batch small enough to be a tile,
+        // though -- applying it to a whole chunk reserves three tokens per span where english
+        // produces about one (2.7 MB of buffer for 0.94 MB of tokens), and that over-allocation cost
+        // 3-5% on the models with no tiled emit, which are exactly the ones that see whole chunks.
+        const TILE_LIKE: usize = 8192;
+        output.reserve(if spans.len() <= TILE_LIKE {
+            spans.len() * MAX_INLINE_IDS + MAX_INLINE_IDS
+        } else {
+            spans.len() + MAX_INLINE_IDS
+        });
         let mut capacity = output.capacity();
         let mut cursor = output.len();
         // A raw write cursor held across the whole chunk. `output.as_mut_ptr()` inside the loop had

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

@@ -94,6 +94,25 @@ impl TryFrom<Sequence> for PipelineSequence {
     }
 }
 
+impl PipelineSequence {
+    /// The single real child, when the `Sequence[Split(regex), ByteLevel]` shape has collapsed to
+    /// one (a byte-map `ByteLevel` converts to `PipelinePreTokenizer::None`, a pure identity pass).
+    /// Same test `pre_tokenize` makes before delegating.
+    pub(crate) fn lone_child(&self) -> Option<&pipeline::PipelinePreTokenizer> {
+        if self.is_deepseek() {
+            return None;
+        }
+        let mut work = self
+            .pre_tokenizers
+            .iter()
+            .filter(|c| !matches!(c, pipeline::PipelinePreTokenizer::None));
+        match (work.next(), work.next()) {
+            (Some(only), None) => Some(only),
+            _ => None,
+        }
+    }
+}
+
 impl pipeline::PreTokenizer for PipelineSequence {
     /// Runs each child in turn, where every child subdivides the spans produced
     /// so far. A child sees only the text of a span (`&text[span]`) and returns

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

@@ -178,6 +178,15 @@ impl PreTokenizer for Split {
     }
 }
 
+impl Split {
+    /// The native grammar this `Split` would route to, if any. Same condition `pre_tokenize` uses,
+    /// named so the pipeline can ask before choosing between the tiled and whole-chunk paths.
+    pub(crate) fn native_grammar(&self) -> Option<crate::utils::Grammar> {
+        self.fsm
+            .filter(|_| !self.invert && self.behavior == SplitDelimiterBehavior::Isolated)
+    }
+}
+
 impl pipeline::PreTokenizer for Split {
     fn pre_tokenize(&self, text: &str, out: &mut Vec<pipeline::Span>) -> Result<()> {
         // A recognized GPT regex in its only real usage — `Isolated`, not inverted — routes