[POC] fast encode: FlatCache + MPHF RankStore + incremental merge#2190
Draft
ArthurZucker wants to merge 142 commits into
Draft
[POC] fast encode: FlatCache + MPHF RankStore + incremental merge#2190ArthurZucker wants to merge 142 commits into
ArthurZucker wants to merge 142 commits into
Conversation
…rst, classify first then FSM simd when possible, pure scalar for known patterns!
Signed-off-by: Luc Georges <luc.sydney.georges@gmail.com>
…it::nfd Expose atomsplit::nfd::nfd_char (per-char NFD via the baked trie + arithmetic Hangul) and use it in BertNormalizer's pipeline strip_accents instead of unicode-normalization's char.nfd(). Drops the last unicode-normalization NFD call from the pipeline normalizer paths (NFD/NFKD/NFC/NFKC + bert all pure-Rust now); the inline Hangul special-case folds into nfd_char. Byte-exact — the bert pipeline-vs-legacy parity tests pass. The legacy NormalizedString paths keep unicode-normalization (they track alignment offsets, which the fast path omits).
…ession) The previous commit folded Bert's Hangul special-case into nfd_char, which then ran is_mark_nonspacing + to_lowercase per jamo — work the old inline path skipped (jamo are caseless Lo, never Mn). Restore the direct S_BASE-arithmetic push for Hangul; non-Hangul chars still decompose via the pure-Rust nfd_char. Byte-exact (bert parity tests pass).
…compose)
The owned decompose/compose path reserved only n + n/16, but decomposition
expands (Hangul 3×, accents ~1.5×), so decompose-heavy input blew through it and
the String realloc'd mid-build (full-buffer copy per grow). bitmap_gen now bakes
{NFD,NFKD}_MAX_EXPAND (max out/in UTF-8 byte ratio; 3 for NFD, 11 for NFKD) and
the kernel reserves input_len * MAX_EXPAND — provably realloc-free. ~2-5% on
decompose-heavy scripts (Hangul biggest), zero cost on the borrow/ASCII path.
Byte-exact; exhaustive + all tk-encode tests pass.
…merge
Fast-encode experiment on top of feat/fast-normalization, kept as a trace (not
for merge). Byte-exact (BPE + sequence unit tests) on gpt2/cl100k/o200k/llama-3.
- flat_cache: thread-local, hash-keyed, L2-sized pretoken->ids cache
(open-addressing slots + bump arena, clear-on-full). The pipeline had none.
- rank_store: MPHF (ptr_hash) pair->(rank,merged) map — one 16B cache-line read
per lookup, ~2x less memory than the AHashMap (llama-3 4.5 vs ~10 MB).
- word/model: incremental rank-array linear merge (O(n) MPHF lookups instead of
O(n^2)) + reused heap scratch; provably 0 hot-path allocations.
- sequence: drop None children so a single-Split Sequence (Llama-3) routes to
the zero-alloc atomsplit FSM instead of the allocating multi-child loop;
reused thread-local scratch for genuine multi-child Sequences.
- examples/xbench_{pipe,iter,alloc}: throughput / per-iteration warm-up /
allocation-count harnesses used for the investigation.
Re-bench vs wordchipper/IREE/bpe-openai/fastokens/tiktoken/tokenizers-main:
fastest single-thread on all 4 models, ~6x near-linear scaling to 8 threads.
|
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. |
…ult) Byte-exact (asserted: windowed path == reference on long input). But it gives NO speedup on cold CJK (gpt2 cmn 54->55, llama-3 29->30, deepseek 27->28) because max_span = 128 for these byte-level vocabs — BPE genuinely builds 128-char tokens (long whitespace / repeat runs), so a provably-safe window is >=256 symbols, which is >= virtually every pre-token. The heap already exploits the same max-length locality (that's why it's O(n log n), not O(n^2)); windowing with a 256-safe-window is O(n * 256), no better. Kept as a trace + a byte-exact test; the real cold-CJK lever remains the rust-gems backtracking encoder.
…filed)
Profiling (stage breakdown + inline(never)-barriered sampling) showed the merge
is 76-95% of encode, and within it: pair-rank lookup ~26-47%, Vec::remove
memmove 8-18%, then heap/scan bookkeeping.
(a) RankStore: drop the redundant splitmix64 — ptr_hash's index_single_part
already FxHashes the u64 key internally (it showed in the profile), so the
extra pass was pure waste. get 40.6->32.7%, FxHash 6.2->3.8% (gpt2/cmn).
(b) merge_linear: replace the Vec::remove (O(n) memmove) with an index-linked
list — a merge is O(1) relinking. _platform_memmove 8.5-17.7% -> 0.
Byte-exact (32 BPE tests incl. the long-input windowed path). Cold (uncached)
throughput: gpt2 +3%, llama-3 eng +7% / cmn +7%, deepseek eng +16% / cmn +8%.
Residual: the fixed-width min-scan now dominates the linear path (traded from
the memmove); a live-index list or the backtracking encoder would cut it next.
Adds examples/xbench_stage + a `profile-noinline` feature for re-profiling.
Restore the pre-token byte-key + verify it on a hash hit before returning the cached ids, so the FlatCache is provably byte-exact rather than trusting 64-bit hash uniqueness (a collision would have served wrong ids, ~2^-49/lookup — not acceptable when exactness is the whole point). A byte mismatch keeps probing (open addressing), so a collision neither returns wrong ids nor a false miss. Measured cost of the verify (warm, single-thread, vs FLATCACHE_HASHONLY=1): gpt2 +2.5..7.0%, cl100k +3.3..5.8%, o200k +1.2..4.2%, llama-3 +1.8..5.5% (mean ~4%). Extra memory: the key arena + 24B slots (~+1.5 MB/thread @ BITS=16). Verify is the default; FLATCACHE_HASHONLY=1 opts into the unsafe fast path only for A/B measurement. 32/32 BPE byte-exact tests pass.
Replace clear-on-full with a generational cull: each entry carries a saturating freq bumped on hit; when the table fills we compact into a scratch buffer keeping only entries reused since the last cull (freq>=1, halved for aging) and dropping the freq-0 one-shots (the unique long pre-tokens that merge once and never recur), then swap buffers — so the hot Zipfian set survives a fill and pollution doesn't. Hot path stays alloc-free (double-buffered); byte-exact verify retained. Measured on a single cold pass over a stream (cache fills as it goes — the real encode-a-document workload, not re-encode). cull vs clear-on-full: - overflow regime (working set > cache, the realistic case): +2..5pp hit rate and up to +15% cold throughput (o200k 6MB multilingual @64k slots: 47 vs 41 MB/s). - an L2-sized cull cache matches a 8x-bigger clear cache (o200k @8k slots cull 74.8%/44 MB/s ~= @64k clear 77.6%/41 MB/s) — same speed, 1/8 the memory. - fits regime (working set < cache): ties on hit rate, cull ~4% slower (freq-bump overhead, no benefit); brutal thrash (148k into 4k): compaction cost > benefit. CACHE_EVICT=clear reverts to wipe-on-full; CACHE_BITS sizes the table.
…ect table Attacks the MPHF rank lookup (16–30% of merge, heaviest on CJK). All byte-exact (32 unit tests + token-stream checksums vs baseline across models/langs/scripts). #3 small-pair direct table — SHIPPED (default on, RANK_NO_SMALL opts out). Byte-byte pairs (a,b<256 — the frequent early merges) go through a 512 KB L2-resident table: no hash, no entry load, no key verify (direct index is exact). Pure-merge throughput +4..10% Latin, +3..16% CJK, every model. #1 versioned heap (HEAP_VER, opt-in) — no-revalidation heap path: each item carries its merged id + a version stamp, so a pop is validated by an array read instead of the baseline's per-pop pair_rank re-lookup. Byte-exact (the version must bump on ANY neighbour change, incl. when the new pair has no merge rule — the subtle case that over-merged CJK until fixed). Benched ~neutral (−1..+6%, ≈+1%): the saved lookup ≈ the fatter 16 B heap entry. #2 fingerprint pre-filter (RANK_FP, opt-in) — 8-bit dense side array to reject absent pairs before the wide entry load. Benched neutral-to-negative: the extra fp load ≈ the saved entry load once the table is ~L2-resident. Kept for reference, off by default. xbench_split: classify-vs-FSM sub-timing; xbench_iter: XBENCH_CKSUM exactness.
Set PtrHashParams::remap = false and query via index_no_remap: the MPHF now returns a slot in [0, max_index()) directly, skipping the Elias-Fano (CachelineEf) remap-to-minimal decode on every lookup (it showed up as its own frame in the profile). The entries table spans the full slot range instead of [0,n) — sparse, ~1.01x larger, some empty slots (the key verify already rejects those, so no behaviour change). Pure-merge cold-off, vs the remap version (both with #3 small-table), median of 3x8: gpt2 +7..9%, cl100k +4..6%, o200k +3..5%, llama3 +4..8%, deepseek +3..6%. Byte-exact: all reference token-stream checksums unchanged across models/langs (incl. Japanese, the heap-path canary). 32/32 unit tests pass.
…h more lightwheight This commit and the next are gonna be cleaning up the stupid AI slope jk
cl100k and o200k ship their pre-tokenizer as Sequence[Split(invert=true, behavior=Removed, <gpt-regex>), ByteLevel] (the tiktoken-conversion convention). The pipeline's Split->FSM fast path only fired for (invert=false, Isolated), so these two — the most important production vocabs — silently fell back to MultiRegex (regex_automata), spending 11-18% of the whole encode in the regex DFA on plain text with zero special tokens. gpt2 (bare ByteLevel, synthesized as Isolated) and deepseek (own FSM path) were unaffected. For a whole-covering GPT regex, (invert=true, Removed) is byte-exactly equivalent to (invert=false, Isolated) — the inverted match set is the gaps, and these patterns leave none. Canonicalize to that form at pipeline build (Split::canonicalized_for_pipeline) when gpt_fsm recognizes the pattern, so cl100k/o200k route to fsm_cl100k / fsm_o200k. Single cold pass, cache on: cl100k +30..63%, o200k +20..39% (gpt2 control ~+3%, noise). regex_automata share 18% -> 0.0%. Byte-exact: all reference token-stream checksums unchanged; bpe + split unit tests pass (2 pre-existing precompiled- normalizer failures are unrelated, fail without this change too).
SBrandeis
reviewed
Jul 15, 2026
Comment on lines
+990
to
+995
| for i in 0..n { | ||
| if alive[i] && ranks[i].0 < best { | ||
| best = ranks[i].0; | ||
| bi = i; | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
One smart trick is to pack rank in high bits with position in low bits (assuming pretoken and the merge cardinality / rank are small enough): rank << 20 | pos
Getting the min also gives you the position in the array for free
Probably negligible vs having another usize anyways
Contributor
|
I made similar changes independently in #2194, if you want to have a look :) (your PR seems to have more scope) Would love to see some bench numbers to see if it speeds up anything |
ArthurZucker
force-pushed
the
feat/fast-normalization
branch
from
July 15, 2026 18:51
ef25e3a to
c016c46
Compare
ArthurZucker
added a commit
that referenced
this pull request
Jul 23, 2026
…33x gap, not merge/split/norm
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.
Draft / experiment — kept as a trace, not for merge.
Fast pipeline-encode POC on top of
feat/fast-normalization. Byte-exact (BPE + sequence unit tests) on gpt2 / cl100k / o200k / llama-3.What's in it
flat_cache.rs— thread-local, hash-keyed, L2-sized pretoken→ids cache (open-addressing + bump arena, clear-on-full). The pipeline had no cache at all.rank_store.rs— MPHF (ptr_hash) pair→(rank, merged) map: one 16-byte cache-line read per lookup, ~2× less memory than theAHashMap(llama-3 4.5 vs ~10 MB).word.rs/model.rs— incremental rank-array linear merge (O(n) MPHF lookups instead of O(n²)) + reused heap scratch; 0 hot-path allocations.sequence.rs— dropNonechildren so a single-SplitSequence(Llama-3) routes to the zero-alloc atomsplit FSM instead of the allocating multi-child loop; reused thread-local scratch for genuine multi-child.examples/xbench_{pipe,iter,alloc}.rs— throughput / per-iteration warm-up / allocation-count harnesses.Re-bench (single thread, mean MB/s, 10 Wikipedia languages)
Multi-thread: ~6× near-linear scaling to 8 threads (gpt2 285, llama-3 307 MB/s), best of the field.
Single-thread numbers are warm-cache and data-dependent (hit-rate varies by language); cold-path merge work keeps low-repeat/CJK fast without the cache (CJK ~2× faster uncached after the incremental merge). All hot paths verified 0 allocs.