feat: multi threaded pipeline by McPatate · Pull Request #2213 · huggingface/tokenizers

McPatate · 2026-07-17T17:24:14Z

Status: ai generated iteration, will trim down and review thoroughly to cut the BS and verify for correctness in more depth

HuggingFaceDocBuilderDev · 2026-07-17T17:27:35Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

new `UnsafeCell` with atomic impl should allow us to scale linearly with core count Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

ArthurZucker

As a reviewer I am missing a lot of help on what is happening at a global scope. So I reviewed smaller scopes and some stuff here and there but we really need to do file split

ArthurZucker · 2026-07-24T04:53:33Z

+    views: Vec<PyStrView>,
+}
+
+// SAFETY: the raw pointers reference CPython's cached UTF-8 buffers, which are


Suggested change

// SAFETY: the raw pointers reference CPython's cached UTF-8 buffers, which are

// SAFETY: the raw pointers reference CPython's cached UTF-8 encoded buffers, which are

not encoded, this is only the input type!

ArthurZucker · 2026-07-24T05:17:42Z

+        // SAFETY: ptr is kept alive by PyStrBatch::_owners, and the underlying buffer is immutable UTF-8
+        unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, len)) }


Suggested change

// SAFETY: ptr is kept alive by PyStrBatch::_owners, and the underlying buffer is immutable UTF-8

unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, len)) }

// SAFETY: the pointer given by view.as_ptr() is non null even for empty str. The memory range is not freed as PyStrBatch::_owners protects it.

let slice =unsafe { std::slice::from_raw_parts(ptr, len)}

// SAFETY: the pointed string IS a utf8 string it does not require another validation (O(len)).

unsafe { std::str::from_utf8_unchecked(slice) };

I think its important to explain why we would go through unsafe when safe can work just as well + not hide to unsafe calls but split them?

Safe would require a clone here, when it's not necessary because python guarantees the buffer won't be gc-ed as long as we have a reference (this is what _owners is for).
Python also guarantees immutability of the buffer, so we can safely create a &str from it. Also the .to_str call we do in PyStrBatch::new guarantees it'll be utf8. So unsafe is needed here, because we're accessing the string's buffer via raw ptr mechanics (we have to if we want to avoid cloning the strings, which is not needed and thanks to this we're 0-copy from python input to rust, we won't be able to achieve something similar without unsafe), but we know the safety guarantees are upheld in this context.

let view = s.to_str()?; // -> guarantees the string buffer's bytes are utf8, then cached and never mutated views.push(PyStrView { ptr: view.as_ptr(), len: view.len(), }); owners.push(s.unbind());

ArthurZucker · 2026-07-24T05:31:09Z

+    tokens.iter().map(|t| t.id).collect()
+}
+
+struct PyStrView {


Suggested change

struct PyStrView {

// A pointer to a utf8-encoded PyString that comes/lives in python world.

struct PyStrView {

PyString is necessarily a python string from python world 😄

Not sure a comment is needed here, imo quite explicit what this is?

ArthurZucker · 2026-07-24T05:42:41Z

+impl IntoInputs for PyStrBatch {
+    type Inputs = PyStrBatch;
+    fn into_inputs(self) -> PyStrBatch {
+        self
+    }
+}
+


as this is gonna change with the updates to support pairs of inputs, makes sense to have just a file for input?

erm, not sure, I think we can extend this struct to support pairs of inputs rather than replacing it. Not sure what you mean by "file for input"

ArthurZucker · 2026-07-24T05:44:05Z

+/// to place it). The consuming thread *assists* the pool while it waits (all
+/// with the GIL released). Dropping the job cancels unclaimed work.


how does it assist? why does it matter?

canceling return claimed and finished work in a wait?

how does it assist? why does it matter?

impl Iterator for EncodeHandle { type Item = (usize, Result<Vec<PipelineToken>>); fn next(&mut self) -> Option<Self::Item> { match &mut self.inner { HandleInner::Ready(it) => it.next(), HandleInner::Streaming { core, state } => state.next_completed(core), } } } impl StreamState { fn next_completed(&mut self, core: &JobCore) -> Option<(usize, Result<Vec<PipelineToken>>)> { if self.next_k >= self.n { return None; } loop { let seq = core.completed_order[self.next_k].load(Ordering::Acquire); if seq != NOT_DONE { self.next_k += 1; return Some((seq, core.take_result(seq))); } // <------ HERE IS THE ASSIST LOGIC ------> if !self.assist_done { if SCRATCH.with(|st| core.run_one(&mut st.borrow_mut())) { continue; } self.assist_done = true; } std::hint::spin_loop(); } } }

I've added a comment to show where the assist logic is.

Why:

sleep is out of the question, too unpredictable and adds too many syscalls

spin (loop { if has_result() { break; } }) burns a core for nothing, although it reacts fast

condvar/futex (park) is the only viable alternative imo, reacts in the µs, but if we do implement it over assist, it has the following downsides:

latency: encode can start straight away, even before the pool has initialised

no context switching, the thread is always alive, no need to pay for µs wake

no need for the workers to signal the main thread that results are ready

overall more complicated, needs quite some code to implement, whereas assist is simply "call the work function we already have each worker do"

canceling return claimed and finished work in a wait?

not sure I understand the question, but if the EncodeHandle is dropped for wtv reason, the ongoing encode jobs are cancelled

ArthurZucker · 2026-07-24T06:16:42Z

+        self.iter().map(|s| (*s).to_owned()).collect()
+    }
+}
+/// One owned encode call's worth of work, shared with the pool workers via


this is also a new file!

ArthurZucker · 2026-07-24T06:20:21Z

+    /// Encode one or many sequences, returning an [`EncodeHandle`]:
+    /// workers keep encoding in the background while the caller
+    /// holds it.
    ///
-    /// todo: wire the post-processing
-    pub fn encode(&self, input: &str, _add_special_tokens: bool) -> Result<Vec<PipelineToken>> {
-        let mut output = Vec::new();
-        let mut pre_tokens = Vec::new();
-        let mut scratch = self.model.init_scratch();
+    /// The job is `'static`: it holds the input storage and a cheap clone of this
+    /// tokenizer handle. Dropping the job cancels unclaimed
+    /// work, and the last worker releases the storage.


this is key / core. Needs a proper doc explaining what is happening or this should be in a readme.

ArthurZucker · 2026-07-24T06:22:59Z

+        assert_eq!(by_ref, serial, "Normalized (borrowed) != serial");
+        let owned = ids(tok.encode(big.clone()).into_single().unwrap());
+        assert_eq!(owned, serial, "Normalized (owned) != serial");
+    }


we need to put the test somewhere else, this file is super hard to read

ArthurZucker · 2026-07-24T06:24:42Z

+//! Process-global worker pool for the parallel encode path: a lazily-built,
+//! library-private `rayon::ThreadPool`.
+//!


why? why do we need this etc!

ArthurZucker · 2026-07-24T06:25:09Z

+pub fn set_num_threads(num_threads: usize) {
+    NUM_THREADS.store(num_threads, Ordering::SeqCst);
+    crate::tokenizer::pool::invalidate();
+}
+
+/// The programmatic worker-count override, if one was set.
+pub(crate) fn num_threads_override() -> Option<usize> {
+    match NUM_THREADS.load(Ordering::SeqCst) {
+        0 => None,
+        n => Some(n),
+    }
+}


mega long due in this lib!

SBrandeis

Partial review will continue later

SBrandeis · 2026-07-24T08:20:02Z

+    ptr: *const u8,
+    len: usize,


Can't we use PyBackedString from pyo3 instead?

It seems to be doing exactly what you've re-implemented here, and provide safe APIs to get &str from the Python-owned object

https://pyo3.rs/main/doc/pyo3/pybacked/struct.pybackedstr

SBrandeis · 2026-07-24T08:34:57Z

+        let job = self.take()?;
+        let res = py.detach(|| job.wait_for_completion());
+        let out = PyResult::from(ToPyResult(res))?;
+        Ok(out.into_iter().map(ids).collect())


Ideally we should be able to send data to python without allocating (maybe through numpy array?)

Probably as a follow-up

yes agree, will tackle later

SBrandeis · 2026-07-24T08:37:36Z

+enum Stage {
+    Frame,
+    Normalize,
+    Split,
+    Model,
+}


Missing the post-process stage (reusing the other const would catch this)

SBrandeis · 2026-07-24T08:38:05Z

+/// Cumulative pipeline stage levels for the ablation ladder, in execution
+/// order: each level runs every stage up to and including itself; `Model` is a
+/// full encode, `Frame` is the special-token scan + iteration only.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum Stage {
+    Frame,
+    Normalize,
+    Split,
+    Model,
+}


The issue with enum is that they cannot be used in const generics

SBrandeis · 2026-07-24T08:38:38Z

+/// ladder). A runtime branch per segment is fine in a bench; the subtraction
+/// cancels its cost.


Why scrapping encode_generic?

SBrandeis · 2026-07-24T08:39:00Z

+    let mut scratch = model.init_scratch();
+    let mut run = |out: &mut Vec<PipelineToken>, pre_tokens: &mut Vec<_>| {
        for chunk in chunks {
            out.clear();
-            let _ =
-                pipeline.encode_generic::<STAGE>(chunk, &mut pre_tokens, &mut scratch, &mut out);
+            for segment in SpecialSegmentIterator::new(chunk, added_vocabulary, false) {
+                match segment {
+                    Segment::SpecialToken(id) => out.push(PipelineToken { id }),
+                    Segment::Text(text) => {
+                        let normalized: Cow<str> = if stage >= Stage::Normalize {
+                            match normalizer {
+                                Some(n) => n.normalize(text).unwrap(),
+                                None => Cow::Borrowed(text),
+                            }
+                        } else {
+                            Cow::Borrowed(text)
+                        };
+                        for seg in SpecialSegmentIterator::new(&normalized, added_vocabulary, true)
+                        {
+                            match seg {
+                                Segment::SpecialToken(id) => out.push(PipelineToken { id }),
+                                Segment::Text(normalized_chunk) => {
+                                    if stage >= Stage::Split {
+                                        pre_tokens.clear();
+                                        pre_tokenizer
+                                            .pre_tokenize(normalized_chunk, pre_tokens)
+                                            .unwrap();
+                                        if stage >= Stage::Model {
+                                            for pre_token in pre_tokens.iter() {
+                                                model
+                                                    .tokenize_pipeline(
+                                                        &normalized_chunk[pre_token.range()],
+                                                        &mut scratch,
+                                                        out,
+                                                    )
+                                                    .unwrap();
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }


Very brittle, needs an update every time we change the encode logic to be faithful

SBrandeis · 2026-07-24T08:45:16Z

+thread_local! {
+    /// Per-thread reusable encode scratch: pool drainers, inline callers and the
+    /// caller-assist path all keep their buffers warm across encode calls
+    /// (reset, never realloc'd) — pool threads persist for the process, so this
+    /// is the cache-warmth contract of [`EncodeState`].
+    static SCRATCH: RefCell<EncodeState> = RefCell::new(EncodeState::new());
+}


Is there a way to get rid of thread_local!?
For example by having each worker own its copy of EncodeState

See #2223 (scratch pool pattern)

thread_local! avoids having to use Arc<Mutex<T>>, it is already owned per thread

it's just not shareable, which is not a problem in our case

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

Brings in the parallel encode runtime from #2213: the `ParallelPlan` ladder, the claim-cursor job core, `EncodeHandle`, and the fork-safe pool. Both branches had rewritten `tokenizer/pipeline.rs`, so #2213's runtime is taken whole and this branch's work is grafted back onto it: - `PipelineInner` carries `normalizers: Vec<_>` (metaspace-as-normalizer) and the wired `PipelinePostProcessor` instead of the single normalizer and the unused `_post_processor`. `stride_boundary` now asks every normalizer; `PipelineNormalizer::Metaspace` answers no, so t5/albert take the `Normalized` plan rather than a raw cut. - The model kernels call the batched `tokenize_spans` instead of looping `tokenize_pipeline` per span -- the point of this branch. - Special-token framing is applied per *input*, where an input's chunks are joined (`take_result` / `encode_serial`), never per chunk. `encode` frames by default; `encode_with` takes the flag. `encode_one` is the single-sequence synchronous path, and `encode_generic::<STAGE>` keeps the bench ablation ladder. - The thread-local `EncodeState` replaces this branch's per-tokenizer `ScratchPool`, but the scratch holds a word cache keyed on word bytes alone -- so it is now rebuilt when the thread's last encode used a *different* tokenizer, not merely a different model kind. Two BPE tokenizers sharing a thread traded ids without this. Also: - `BucketVocabStore::build` returns the empty store instead of panicking on an empty vocab; `BPE::default()` is one, and the plan tests build from it. - #2213's oracle keeps its own file (`pipeline_parallel_oracle.rs`) next to the per-model one, and compares against `add_special_tokens = true` now that the post-processor runs -- which also checks the framing lands once per input. - `bpe_pipeline_oracle` skips a model whose *legacy* reference can't encode without `fancy-regex` (deepseek): with the backend on, all four are exact.

* fix(pipeline): don't stride an input that has no cut points `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. * 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. * Delete tokenizers/tk-encode/examples/ab_giga_mt.rs

McPatate requested review from ArthurZucker and SBrandeis July 17, 2026 17:24

McPatate added 3 commits July 20, 2026 11:41

feat: impl pipeline multi-threading

c6d306f

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

feat: introduce parallelplan strategy

22033b3

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

refactor: remove unused add_special_tokens

61a82cd

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

McPatate force-pushed the feat/multi_threaded_pipeline branch from b5c7c18 to 61a82cd Compare July 20, 2026 09:41

McPatate added 11 commits July 20, 2026 16:28

chore: cleanup pipeline bindings

f866644

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: cleanup pipeline benchmark

addabaa

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

fix: update encode API

3765e0a

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

feat: add special tokens & normalization chunking

1c0f9e5

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: cleanup

b644aba

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: more cleanup

8581dfa

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: harmonize naming for segments and parallel splitting

8aec88b

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: docs conciseness attempt

e52925c

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

perf: reduce allocations + avoid false sharing

2f0e387

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

refactor: remove channel because of immense contention

502347f

new `UnsafeCell` with atomic impl should allow us to scale linearly with core count Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

chore: trim docs

8596486

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

ArthurZucker reviewed Jul 24, 2026

View reviewed changes

SBrandeis reviewed Jul 24, 2026

View reviewed changes

chore: address comments

852c82a

Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>

ArthurZucker mentioned this pull request Aug 4, 2026

perf(pipeline): don't stride an input that has no cut points #2283

Merged

ArthurZucker added the target label Aug 6, 2026

	// SAFETY: the raw pointers reference CPython's cached UTF-8 buffers, which are
	// SAFETY: the raw pointers reference CPython's cached UTF-8 encoded buffers, which are

		// SAFETY: ptr is kept alive by PyStrBatch::_owners, and the underlying buffer is immutable UTF-8
		unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, len)) }

-        // SAFETY: ptr is kept alive by PyStrBatch::_owners, and the underlying buffer is immutable UTF-8
-        unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, len)) }
+        // SAFETY: the pointer given by view.as_ptr() is non null even for empty str. The memory range is not freed as PyStrBatch::_owners protects it.
+        let slice =unsafe { std::slice::from_raw_parts(ptr, len)}
+        // SAFETY: the pointed string IS a utf8 string it does not require another validation (O(len)).
+        unsafe { std::str::from_utf8_unchecked(slice) };

	struct PyStrView {
	// A pointer to a utf8-encoded PyString that comes/lives in python world.
	struct PyStrView {

		/// to place it). The consuming thread assists the pool while it waits (all
		/// with the GIL released). Dropping the job cancels unclaimed work.

		/// ladder). A runtime branch per segment is fine in a bench; the subtraction
		/// cancels its cost.

Uh oh!

Conversation

McPatate commented Jul 17, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jul 17, 2026

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SBrandeis left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants