From f9342ae2c9de8c98eef5f33171b1b2cfbabd2e96 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 5 Aug 2026 00:01:13 +0900 Subject: [PATCH 1/3] fix(pipeline): don't stride an input that has no cut points MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `plan()` chose `Raw` whenever the *config* exposed a stride boundary, without asking whether the *input* contains one. Punctuation-terminated CJK contains none at all: Chinese prose has no space anywhere, and a newline after `。` is not a legal cut because the punct rule ` ?[…]+[\r\n]*` absorbs it into the punct token (`NEWLINE_PREV` excludes punctuation for exactly that reason). `data/corpora/chinese.txt` is 100% such newlines -- 722 of them, every one preceded by `。`, and zero spaces. Striding then did what `stride_range`'s own doc comment warns about: "text with no boundaries at all degrades to stride 0 owning the whole input". Measured on a 4 MB doc: 513 strides tiled, **1 resolved, covering all 4199874 bytes**, after scanning the document twice looking for cuts that do not exist. 0.28x the plain serial encode, and flat across thread counts because there is only ever one work unit. `plan()` now probes the longest input for an actual cut (bounded at 32 KB, once per encode) and falls through to `Pretokenized` when there is none -- which pre-tokenizes serially and parallelises the model over span groups. That prefix runs the normal pre-tokenizer, so it is `bitsplit` wherever bitsplit is wired (gpt2, cl100k), which is what makes it cheap enough to prefer over a degenerate stride. gpt2 chinese, 14 x 4 MB docs, MB/s ours/giga -- was 1193/1243 (0.94x), now 3004/1120 (2.68x); on a cool box the same fix measured 4798/1243 (3.86x). Single thread on one doc: 260 -> 732. Token counts unchanged (23531424), and every plan was already byte-exact, so plan choice only ever trades throughput. Spaced text still probes as cuttable and takes `Raw` unchanged. `cut_exists` passes `lo = 1`, not 0: `boundary_in_window` reads one byte of left context via `block_lo - 1`, so 0 underflows to `usize::MAX` and spins forever in the char-boundary walk-back (it hangs, it does not panic, in release). Byte 0 is never a cut anyway. Written with nested `if let` rather than let-chains: this branch is edition 2018. --- .../tk-encode/src/tokenizer/pipeline.rs | 127 +++++++++++++++--- 1 file changed, 108 insertions(+), 19 deletions(-) diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index eb647501b..58f3ed248 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -521,6 +521,27 @@ fn boundary_fsm(window_tags: &[u8]) -> Option { (1..window_tags.len()).find(|&i| safe_fsm_cut(prev_char_tag(window_tags, i), window_tags[i])) } +/// Whether `text` contains any `boundary` cut at all, deciding within a bounded prefix. +/// +/// Used by [`PipelineTokenizer::plan`] to reject a cut-based plan for input that has no cuts. +/// Only the first few strides are examined: a text whose opening `PROBE` bytes hold no cut is +/// treated as uncuttable, which is right for the case this exists for (Chinese/Japanese prose +/// has none anywhere) and merely gives up some parallelism in the pathological case of a +/// document that is boundary-free only at the front. Bounded so the probe stays O(1) in the +/// input: it runs once per `encode`, before any work is dispatched. +fn cut_exists(text: &str, boundary: StrideBoundary) -> bool { + const PROBE: usize = 4 * PARALLEL_MIN_TOTAL_BYTES; + // `boundary_in_window` reads one byte of left context (`block_lo - 1`), so `lo` must be + // at least 1 -- passing 0 underflows and hangs the char-boundary walk-back. Byte 0 is + // never a cut anyway: stride 0 starts at 0 unconditionally. + if text.len() < 2 { + return false; + } + let hi = text.len().min(PROBE); + let hi = (1..=hi).rev().find(|&i| text.is_char_boundary(i)).unwrap_or(1); + boundary_in_window(text, 1, hi, boundary).is_some() +} + /// Whether `token`'s own bytes contain a cut of `boundary` strictly inside it — /// i.e. a stride could split this token and each half would mis-frame it. Used /// by `normalized_added_token_blocks_stride` to disqualify striding for a `normalized` @@ -1316,7 +1337,9 @@ impl PipelineTokenizer { let mut spans: Vec = Vec::new(); let mut preresolved: Vec<(usize, usize, u32)> = Vec::new(); - let plan = self.plan(); + // Probed on the longest input: it dominates the wall time, and a short one is + // unrepresentative (a 20-byte first input would decide the plan for a 4 MB second). + let plan = self.plan(refs.iter().copied().max_by_key(|s| s.len()).unwrap_or("")); // `Normalized`/`Pretokenized` earn their serial prefix only for a // segment larger than a fair per-thread share of the batch. Below that, // batch parallelism balances it as one whole unit; the prefix would just @@ -1613,21 +1636,43 @@ impl PipelineTokenizer { /// How this config splits each special-free segment (specials are peeled /// first). The cheapest safe split wins. - fn plan(&self) -> ParallelPlan { + /// + /// `sample` is the text about to be encoded. A boundary the *config* allows is not the + /// same as a boundary the *input* contains, and the difference is not academic: + /// punctuation-terminated CJK has no cut at all (no spaces anywhere, and + /// [`NEWLINE_PREV`] rejects a newline after punctuation because the punct rule + /// ` ?[…]+[\r\n]*` absorbs it). Striding such an input scans it twice for cuts that do + /// not exist and then hands the whole thing to one worker — measured at 0.28x the plain + /// serial encode, and it does not parallelise at any thread count. So probe the input + /// before committing to a cut-based plan. + /// + /// Every plan is byte-exact (the oracles pin all of them), so this only ever trades + /// throughput — a mis-probe cannot change the ids. + fn plan(&self, sample: &str) -> ParallelPlan { if let Some(boundary) = self.stride_boundary() { - // `Raw`: no serial prefix, so always worth it when available. - ParallelPlan::Raw(boundary) - } else if matches!(self.inner.model, PipelineModel::WordLevel(_)) { + if cut_exists(sample, boundary) { + // `Raw`: no serial prefix, so always worth it when the input can be cut. + return ParallelPlan::Raw(boundary); + } + } + if matches!(self.inner.model, PipelineModel::WordLevel(_)) { // WordLevel's model + pre-tokenize are cheap, so a serial prefix // would dwarf the parallel part. Lean on batch parallelism instead. - ParallelPlan::Whole - } else if let Some(boundary) = self.normalized_stride_boundary() { + return ParallelPlan::Whole; + } + if let Some(boundary) = self.normalized_stride_boundary() { // Normalizer rules out a raw cut, but the pre-tokenizer can cut the - // (serially) normalized text. - ParallelPlan::Normalized(boundary) - } else { - ParallelPlan::Pretokenized + // (serially) normalized text. Probed on the raw sample: normalization does not + // manufacture whitespace, so a text with no cut raw has none normalized either. + if cut_exists(sample, boundary) { + return ParallelPlan::Normalized(boundary); + } } + // No usable cut: pre-tokenize serially and parallelise the model over span groups. + // The serial prefix goes through the normal pre-tokenizer, so it is `bitsplit` + // wherever `bitsplit` is wired (gpt2, cl100k) -- which is what makes this fallback + // cheap enough to prefer over a degenerate stride. + ParallelPlan::Pretokenized } /// The pre-tokenizer's cut boundary for **already-normalized** text, gated @@ -2079,6 +2124,11 @@ impl ModelScratch for PipelineModelScratch {} #[cfg(test)] mod tests { use super::*; + + /// A `plan()` probe input that plainly contains cuts (spaces after letters), so these + /// tests keep asserting what the *config* allows rather than what a sample happens to + /// hold -- `cut_exists` is covered separately by `uncuttable_input_skips_striding`. + const CUTTABLE: &str = "the quick brown fox jumps over the lazy dog and runs away fast "; use crate::models::bpe::BPE; use crate::models::wordpiece::WordPiece; use crate::pre_tokenizers::byte_level::ByteLevel; @@ -2314,6 +2364,45 @@ mod tests { /// added token disables raw cutting only when it is `normalized` (matched /// after normalization, so invisible to the special split's raw pass); a raw/special one /// is peeled first, so a stride can never bisect it. + /// The bug this guards: a config CAN cut (gpt2 exposes `boundary_fsm`) but the INPUT + /// cannot. Chinese prose has no space anywhere and ends its lines with `。`, and a + /// newline after punctuation is not a legal cut (the punct rule absorbs it), so striding + /// found zero cuts, gave one worker the whole document, and scanned the document twice + /// on the way. Measured at 0.28x the serial encode and flat across thread counts. + #[test] + fn uncuttable_input_skips_striding() { + let split = SplitPretok::new( + SplitPattern::Regex(GPT2_REGEX_STR.to_owned()), + SplitDelimiterBehavior::Isolated, + false, + ) + .unwrap(); + let mut tok = Tokenizer::new(crate::models::bpe::BPE::default()); + tok.with_pre_tokenizer(Some(split)); + let pipe = PipelineTokenizer::try_from(&tok).unwrap(); + + // The config itself allows a cut -- that is why the old code chose `Raw`. + assert!(pipe.stride_boundary().is_some()); + + // Punctuation-terminated Chinese: no space, and every newline follows `。`. + let cjk = "汉字汉字汉字汉字汉字。\n".repeat(2000); + assert!( + !cut_exists(&cjk, pipe.stride_boundary().unwrap()), + "`。` then newline must not count as a cut" + ); + assert!( + matches!(pipe.plan(&cjk), ParallelPlan::Pretokenized), + "uncuttable input must not be strided" + ); + + // A newline after a Han *letter* is a legal cut, so that input still strides. + let cuttable = "汉字汉字汉字汉字汉字\n".repeat(2000); + assert!(matches!(pipe.plan(&cuttable), ParallelPlan::Raw(_))); + + // And ordinary spaced text is unaffected. + assert!(matches!(pipe.plan(CUTTABLE), ParallelPlan::Raw(_))); + } + #[test] fn space_run_gating() { use crate::AddedToken; @@ -2389,7 +2478,7 @@ mod tests { .unwrap(), )); let tok = PipelineTokenizer::try_from(&tok).unwrap(); - assert!(matches!(tok.plan(), ParallelPlan::Raw(_))); + assert!(matches!(tok.plan(CUTTABLE), ParallelPlan::Raw(_))); let big = "aa bb,cc! aa\tbb cc\n\n".repeat(2000); // ~44 KB, mixed runs let chunks = (0..big.len().div_ceil(PARALLEL_MIN_TOTAL_BYTES)) @@ -2499,7 +2588,7 @@ mod tests { .unwrap(); let pipe = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(pipe.plan(), ParallelPlan::Raw(_)), + matches!(pipe.plan(CUTTABLE), ParallelPlan::Raw(_)), "raw affix token must keep Raw" ); // Two ~20 KB segments around one lstrip token → both stride. @@ -2521,7 +2610,7 @@ mod tests { let oracle = Tokenizer::from_file("../data/llama-3-tokenizer.json").unwrap(); let tok = PipelineTokenizer::try_from(&oracle).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Raw(_)), + matches!(tok.plan(CUTTABLE), ParallelPlan::Raw(_)), "llama-3's raw-only specials must not disqualify Raw" ); } @@ -2564,7 +2653,7 @@ mod tests { #[test] fn split_at_model_matches_serial() { let tok = split_at_model_pipeline(false); - assert!(matches!(tok.plan(), ParallelPlan::Pretokenized)); + assert!(matches!(tok.plan(CUTTABLE), ParallelPlan::Pretokenized)); let big = "aa.bb.cc.".repeat(3000); // ~27 KB, punctuation-delimited let pretokenized = tok.pretokenize_segment(&big, &big).unwrap(); @@ -2590,7 +2679,7 @@ mod tests { #[test] fn split_at_model_normalizer_and_specials_match_serial() { let tok = split_at_model_pipeline(true); - assert!(matches!(tok.plan(), ParallelPlan::Pretokenized)); + assert!(matches!(tok.plan(CUTTABLE), ParallelPlan::Pretokenized)); let half = "AA.BB.cc.".repeat(2000); let big = format!("{half}{}", "CC.aa.BB.".repeat(2000)); @@ -2637,7 +2726,7 @@ mod tests { ); // WordLevel model: the plan cannot escalate to Pretokenized either. assert!( - matches!(pipe.plan(), ParallelPlan::Whole), + matches!(pipe.plan(CUTTABLE), ParallelPlan::Whole), "Strip + WordLevel must fall to batch-level parallelism only", ); } @@ -2671,7 +2760,7 @@ mod tests { .unwrap(); let tok = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Whole), + matches!(tok.plan(CUTTABLE), ParallelPlan::Whole), "Strip + WordLevel plan must be Whole (the special split does the splitting)" ); @@ -2712,7 +2801,7 @@ mod tests { .unwrap(); let tok = PipelineTokenizer::try_from(&tok).unwrap(); assert!( - matches!(tok.plan(), ParallelPlan::Normalized(_)), + matches!(tok.plan(CUTTABLE), ParallelPlan::Normalized(_)), "Prepend (unsafe normalizer) + WhitespaceSplit + WordPiece must pick `Normalized`" ); let big = "aa bb cc\n".repeat(4000); // ~36 KB, forces striding From 04ff7e4e7dc4738202d9c3db60528927738a87ad Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 5 Aug 2026 00:01:13 +0900 Subject: [PATCH 2/3] bench(ab_giga_mt): multi-thread arm of the gigatoken A/B One 4 MB document per thread through the parallel `encode`, matching gigatoken's `hf_mt` (`encode_docs_ragged` over the same docs) in thread count, work per thread and MiB/s. 14 threads, best-of-3 interleaved, ours/giga: gpt2 english 6240/9229 .68x | code 4252/3056 1.39x | chinese 1181/1251 .94x | russian 4463/1185 3.77x llama-3 english 5075/8766 .58x | code 3380/4071 .83x | chinese 1380/1926 .72x | russian 4894/1638 2.99x Geomean: gpt2 1.35x, llama-3 1.01x. Scaling against each side's own 1-thread encode is the more useful read: english 5.4x/4.5x (giga 6.2x/6.5x), code 6.8x/4.9x (3.9x/5.2x), russian 6.4x/4.9x (2.0x/2.6x), but **chinese only 1.6x on both models** (giga 2.5x/3.8x) -- the whitespace-boundary striding limit already noted for single-document CJK, which is why gigatoken overtakes us at 14 threads on chinese despite our 1.5-1.7x single-thread lead there. --- tokenizers/tk-encode/examples/ab_giga_mt.rs | 55 +++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 tokenizers/tk-encode/examples/ab_giga_mt.rs diff --git a/tokenizers/tk-encode/examples/ab_giga_mt.rs b/tokenizers/tk-encode/examples/ab_giga_mt.rs new file mode 100644 index 000000000..5b4e8a24d --- /dev/null +++ b/tokenizers/tk-encode/examples/ab_giga_mt.rs @@ -0,0 +1,55 @@ +//! Our side of the gigatoken multi-thread A/B: one 4 MB document per thread through the +//! parallel `encode`, matching gigatoken's `hf_mt` (`encode_docs_ragged` over the same +//! docs) in thread count, work per thread and MiB/s convention. +//! +//! AB_TOKENIZER= AB_CORPUS= AB_THREADS= AB_PASSES= +//! +//! Prints the best pass, so the caches are warm -- the same face `hf_mt` reports. +use std::convert::TryFrom; +use std::hint::black_box; +use std::path::PathBuf; +use std::time::Instant; + +use tk_encode::Tokenizer; +use tk_encode::pipeline::PipelineTokenizer; + +const DOC: usize = 4 * 1024 * 1024; + +fn env_usize(key: &str, default: usize) -> usize { + std::env::var(key) + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(default) +} + +fn main() { + let tok = PathBuf::from(std::env::var("AB_TOKENIZER").expect("AB_TOKENIZER")); + let corpus = PathBuf::from(std::env::var("AB_CORPUS").expect("AB_CORPUS")); + let threads = env_usize("AB_THREADS", 1); + let passes = env_usize("AB_PASSES", 4); + // Our pool reads this; set it before the first encode so the pool is built at this size. + tk_encode::utils::parallelism::set_num_threads(threads); + + let legacy = Tokenizer::from_file(&tok).expect("load"); + let pipe = PipelineTokenizer::try_from(&legacy).expect("pipeline"); + + let text = std::fs::read_to_string(&corpus).expect("corpus"); + let one = text.repeat(DOC.div_ceil(text.len())); + let docs: Vec<&str> = (0..threads).map(|_| one.as_str()).collect(); + let total = one.len() * threads; + + let mut best = 0.0f64; + let mut tokens = 0usize; + for _ in 0..passes { + let start = Instant::now(); + let out = pipe.encode(&docs[..]).wait_for_completion().expect("encode"); + let mbs = total as f64 / start.elapsed().as_secs_f64() / (1024.0 * 1024.0); + tokens = out.iter().map(Vec::len).sum(); + black_box(out); + best = best.max(mbs); + } + println!( + "{:<10} {threads:>2}t {best:>7.0} MB/s {tokens} tokens", + corpus.file_stem().unwrap().to_string_lossy() + ); +} From 4b5ad13b15b70ea6f03c971598b347785a429ad4 Mon Sep 17 00:00:00 2001 From: Arthur <48595927+ArthurZucker@users.noreply.github.com> Date: Wed, 5 Aug 2026 03:51:22 +0200 Subject: [PATCH 3/3] Delete tokenizers/tk-encode/examples/ab_giga_mt.rs --- tokenizers/tk-encode/examples/ab_giga_mt.rs | 55 --------------------- 1 file changed, 55 deletions(-) delete mode 100644 tokenizers/tk-encode/examples/ab_giga_mt.rs diff --git a/tokenizers/tk-encode/examples/ab_giga_mt.rs b/tokenizers/tk-encode/examples/ab_giga_mt.rs deleted file mode 100644 index 5b4e8a24d..000000000 --- a/tokenizers/tk-encode/examples/ab_giga_mt.rs +++ /dev/null @@ -1,55 +0,0 @@ -//! Our side of the gigatoken multi-thread A/B: one 4 MB document per thread through the -//! parallel `encode`, matching gigatoken's `hf_mt` (`encode_docs_ragged` over the same -//! docs) in thread count, work per thread and MiB/s convention. -//! -//! AB_TOKENIZER= AB_CORPUS= AB_THREADS= AB_PASSES= -//! -//! Prints the best pass, so the caches are warm -- the same face `hf_mt` reports. -use std::convert::TryFrom; -use std::hint::black_box; -use std::path::PathBuf; -use std::time::Instant; - -use tk_encode::Tokenizer; -use tk_encode::pipeline::PipelineTokenizer; - -const DOC: usize = 4 * 1024 * 1024; - -fn env_usize(key: &str, default: usize) -> usize { - std::env::var(key) - .ok() - .and_then(|v| v.trim().parse().ok()) - .unwrap_or(default) -} - -fn main() { - let tok = PathBuf::from(std::env::var("AB_TOKENIZER").expect("AB_TOKENIZER")); - let corpus = PathBuf::from(std::env::var("AB_CORPUS").expect("AB_CORPUS")); - let threads = env_usize("AB_THREADS", 1); - let passes = env_usize("AB_PASSES", 4); - // Our pool reads this; set it before the first encode so the pool is built at this size. - tk_encode::utils::parallelism::set_num_threads(threads); - - let legacy = Tokenizer::from_file(&tok).expect("load"); - let pipe = PipelineTokenizer::try_from(&legacy).expect("pipeline"); - - let text = std::fs::read_to_string(&corpus).expect("corpus"); - let one = text.repeat(DOC.div_ceil(text.len())); - let docs: Vec<&str> = (0..threads).map(|_| one.as_str()).collect(); - let total = one.len() * threads; - - let mut best = 0.0f64; - let mut tokens = 0usize; - for _ in 0..passes { - let start = Instant::now(); - let out = pipe.encode(&docs[..]).wait_for_completion().expect("encode"); - let mbs = total as f64 / start.elapsed().as_secs_f64() / (1024.0 * 1024.0); - tokens = out.iter().map(Vec::len).sum(); - black_box(out); - best = best.max(mbs); - } - println!( - "{:<10} {threads:>2}t {best:>7.0} MB/s {tokens} tokens", - corpus.file_stem().unwrap().to_string_lossy() - ); -}