poc(encode): target encode path — bitsplit + batched model + fused cache probe - #2279
Open
ArthurZucker wants to merge 131 commits into
Open
poc(encode): target encode path — bitsplit + batched model + fused cache probe#2279ArthurZucker wants to merge 131 commits into
ArthurZucker wants to merge 131 commits into
Conversation
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
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>
…ed cache probe
Everything the split and the model stages have been converging on, in one tree, so
the end state can be measured rather than argued about. Not for merge as-is: the
bitsplit crate is vendored here and belongs in its own library.
- pre-tokenization runs as a bitstream program (`bitsplit`) for the gpt2 and
cl100k regexes: 64 bytes decided per register op instead of one token per
unpredictable branch. split 0.64 -> 0.34 ns/B on English.
- the model takes a whole chunk of pre-tokens at a time (#2278)
- a cache hit is one load and an unconditional store of the slot's lanes,
written straight at the caller's cursor (#2278)
- the word key is one page-guarded load and a masked shift (#2276), folded
with hardware CRC (#2277)
Full encode, 4 MB inputs, single thread, warm cache, aarch64, vs gigatoken 0.10
on the same machine and the same corpora:
gpt2 english 1228/1554 code 674/806 chinese 970/517 russian 847/632
llama-3 english 1265/1405 code 749/802 chinese 998/537 russian 1022/664
Geomean against gigatoken: gpt2 1.14x, llama-3 1.24x.
|
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. |
bitsplit's block builder is NEON-only; everywhere else it falls back to the portable byte-at-a-time reference, which is slower than the FSM it replaces (split 0.89 vs 0.64 ns/B on English). The CI benchmark runs on x86_64, so the routing was handing it a slower splitter. `bitsplit::FAST_BUILDER` says whether the fast build exists; the GPT arms take it only when it does, so x86 keeps the FSM until the x86 transpose lands. Also: clippy (needless_range_loop, needless_return, byte groupings, redundant closure) and rustfmt.
The transpose was NEON-only, so every other target fell back to the portable byte-at-a-time builder -- slower than the FSM it replaces, which is what the x86_64 CI benchmark was measuring. x86 has the easier half of this: `movemask` is native, so a bit-plane is `movemask(slli_epi16(code, 7 - bit))` -- shifting 16-bit lanes lands bit k of both bytes exactly where movemask reads. The LUT is the harder half: `pshufb` indexes 16 entries and tags run to 0x26, so it is one shuffle per high nibble with a select between them. Byte-exact: `verify` built for x86_64-apple-darwin and run under Rosetta (which reports SSSE3, so the SIMD path is the one taken) passes the whole suite -- 220k fuzz cases, the exhaustive 3- and 4-char sweeps, all three grammars, ten corpora re-sliced at 70 offsets each. `fast_builder()` is now a runtime check on x86 (SSSE3 detection) rather than a target const, and the GPT routing keeps its FSM path when it returns false.
At 10 kB a large share of each pass is per-call overhead -- scratch checkout, added-token scan, post-process -- which rates implementations for work that is not tokenizing. The throughput phase now encodes one ~10 MB input per fixture, repeating a fixture that is shorter than that. The scaling sweep parallelises across chunks, so it re-slices that same input into ~512 kB pieces rather than encoding more of them: every thread still has units to take, and the phase costs exactly what it did before. The memory child keeps its 10 kB prefix -- its `rss0` is supposed to measure the tokenizer, not a corpus sitting next to it. Note this makes phase 1 encode 10x the bytes it used to; the released-crate baseline is the slow side of that.
Brings in the BPE merge work from #2241 and #2275: the multipass and two-tier-queue engines, the rank tables and the pretoken->rank conversion. `models/bpe/model.rs` is gone -- #2241 split it into `bpe_model.rs`, `bpe_scratch.rs` and the merge engines -- so the word cache and the batched `tokenize_spans` this branch added to it are re-wired onto the new engine: `BpeScratch` carries `word_cache` again, and both pipeline entry points emit through `tables.unmap` instead of `Word::get_chars_iter`.
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.
The merge had four ways to silently drop parallelism (plan selection, the scratch identity check, the framing hook, the pool), and none of them show up as a test failure. `par_check` prints serial vs one-document-parallel vs batch-parallel MB/s per model, so a regression is one command away. Also clears the two clippy warnings the merged tree inherited.
…gh encode_generic `ab_giga` mirrors gigatoken's own `encode_st` protocol (one repeated buffer, pass 0 cold / later passes warm, best-of over the warm passes, token count printed) so the two can be compared on the same inputs. Token counts match gigatoken exactly on all 8 model x corpus pairs, so the numbers are apples-to-apples. Single-thread warm, 4 MB, interleaved best-of-3 (this box swings ~20% by machine state, so runs must be interleaved -- a single run cannot resolve 10%): gpt2 english 1145/1488 .77x | code 619/776 .80x | chinese 776/493 1.57x | russian 698/602 1.16x llama-3 english 1139/1354 .84x | code 714/780 .92x | chinese 953/503 1.89x | russian 1001/627 1.60x Geomean vs gigatoken: gpt2 1.03x, llama-3 1.24x. `encode_one` now goes through the fused `encode_generic` instead of the worker kernel `encode_one_with`, which splits the pipeline across two functions so a worker can enter at either. Worth ~3% on a whole sequence; it is not the whole story (see below). Measured cost of the merge itself: same harness on the pre-merge tree is 1-8% faster on 7 of 8 pairs (~3% mean). That is the `Arc<PipelineInner>` pointer hop the parallel runtime needs for its 'static workers, and it buys 4-10x on multiple threads.
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.
The harness returned a fresh `Vec` per encode (`encode_one`) while gigatoken's reserves once
and `clear()`s per pass. That cost an allocation plus a first touch of the whole token array
every pass -- and it scales with TOKENS, not bytes, so it quietly penalised exactly the
token-dense corpora: chinese emits 2.9 M tokens for a 4 MB input against english's 0.8 M.
Chinese was reading 1.52x when it is 1.87x. Corrected, single thread, warm, ours/giga:
gpt2 english .81x | code .83x | dense 1.02x | russian 1.36x | greek 1.35x
arabic 1.33x | hindi 1.26x | thai 1.33x | korean 1.49x | chinese 1.87x geomean 1.23x
llama-3 english .91x | code .95x | dense .99x | russian 1.62x | greek 1.59x
arabic 1.72x | hindi 1.59x | thai 1.70x | korean 1.84x | chinese 1.91x geomean 1.43x
Token counts match gigatoken exactly on all 20 pairs. Also sweeps every corpus now, not four,
and takes deepseek-v4 when present (0.55-0.92x there -- that model wants the char-fold, which
is not in this tree).
`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.
SBrandeis
force-pushed
the
feat/bpe-cache-min
branch
3 times, most recently
from
August 4, 2026 23:14
8195f43 to
4b47341
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Supersedes #2272. This is the same intent — one tree with the optimization work in it — but
built on the current cache stack and measured end to end against gigatoken on this machine.
Not for merge as-is:
bitsplitis vendored here astokenizers/bitsplit. It belongs in its ownlibrary; this PR exists so the end state can be measured instead of argued about, and so the
individual PRs underneath it can be reviewed knowing what they add up to.
What is in it
memcpyMeasured
Full encode, 4 MB inputs, single thread, warm cache, aarch64,
lto = "fat". gigatoken 0.10 builtand run on the same machine, same corpora, same tokenizer.json — token counts match it exactly on
all three models, so this is like for like.
Geomean vs gigatoken: gpt2 1.14×, llama-3 1.24×. Ahead overall; within 7–10% on its two best
cells, well ahead on everything with long pre-tokens.
For scale,
feat/bpe-cache-minon the same bench is english 368 / code 178, so this is 3.3× thebase on English and 3.8× on code.
Reproduce with
cargo run --release --example stage_bench(included), which prints the marginalns/B of each stage.
On the CI benchmark
The first push of this branch regressed it: the transpose was NEON-only, the CI runner is x86_64,
so the routing handed it bitsplit's byte-at-a-time fallback (split 0.89 ns/B) in place of the FSM
(0.64). Fixed twice over — there is now an SSSE3 builder, and
bitsplit::fast_builder()gates therouting so any target without a SIMD build keeps the FSM.
Where the remaining gap is
With the model stubbed to answer every pre-token with one token, the per-pre-token call structure
measured 0.64 ns/B before #2278 — gigatoken's entire encode — and is 0.09 now. What is left
on English is split 0.33 + model 0.47; matching gigatoken outright needs the split cheaper still,
which is bitsplit's own roadmap (o200k and the class-run family have no bitstream grammar yet, so
atomsplit::fsmcannot be retired).Dead ends, measured, so nobody repeats them
batched. The table is L1-resident — 512 slots performs the same as 262144.
disable it and English collapses 1554 → 32 MB/s. Our merge is 1.8–6.0× faster than theirs uncached.
PipelineTokenizer benchmark
9 / 10 models supported — PipelineTokenizer vs
tokenizersv0.23.1 (latest release) · ~10 kB inputs · add_special_tokens on · single thread + 1/2/4/8/max-thread sweep86bd5faa7 · 2026-08-04 11:15 UTC· Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · 32 coresvs base branch (
e87a72c56) — per-model geomean ×speedup of this PR's PipelineTokenizer against the base branch's; regressions in red.Decode
Round-trip: v0.23.1
encode_fastproduces the id streams (same fixtures,add_special_tokens=true); both implementations decode those SAME ids withskip_special_tokens=false. MB/s counts decoded text bytes.bert-base-uncased — normalizer-heavy WordPiece · ×8.72 vs v0.23.1 · ×1.11 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 12+0 (peak 12) · Pipeline 8+2 (peak 17)
deepseek-v4 — deepseek 3-regex split-heavy byte-level BPE · ×25.03 vs v0.23.1 · ×3.99 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 62+0 (peak 68) · Pipeline 82+0 (peak 82)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).gemma-4 — byte-fallback BPE, Metaspace-style split (gemma-4) · ×1.85 vs v0.23.1 · ×0.41 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 304+0 (peak 371) · Pipeline 275+0 (peak 371)
gpt2 — gpt2 ByteLevel regex · ×35.09 vs v0.23.1 · ×3.03 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 25+2 (peak 27) · Pipeline 27+0 (peak 28)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).gpt-oss — o200k-regex byte-level BPE (gpt-oss) · ×20.66 vs v0.23.1 · ×2.74 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 241+0 (peak 315) · Pipeline 234+0 (peak 316)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).glm-5.2 — cl100k-variant regex byte-level BPE (glm-5.2) · ×24.40 vs v0.23.1 · ×2.70 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 169+0 (peak 231) · Pipeline 170+0 (peak 231)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).llama-2 — model-bounded BPE, no pre-tokenizer · ×4.66 vs v0.23.1 · ×0.59 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 16+0 (peak 20) · Pipeline 24+0 (peak 24)
llama-3 — cl100k-regex byte-level BPE (llama-3), single regex · ×29.54 vs v0.23.1 · ×2.79 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 73+0 (peak 95) · Pipeline 92+0 (peak 95)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).mistral-small-4 — tekken byte-level BPE, 1k added specials (mistral-small-4) · ×9.15 vs v0.23.1 · ×1.01 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 152+0 (peak 193) · Pipeline 109+0 (peak 195)
Pre-tokenize:
classify + fsmvs regex engines — ns/byte, lower better. The fsm is the scalar jump-table in both pipe columns; SIMD / scalar is the classify pass (regex pre-tokenizers have no SIMD fsm).×vs= engine ÷ our pipeline (SIMD / scalar classify);onig&pcre2(JIT) are C,fancyis pure-Rust fancy-regex,logosis a compile-time DFA lexer (approximate grammar; n/a for deepseek).Not yet supported:
t5-base