perf(tk-encode): SIMD-first encode pipeline — fused split cache + de-virtualized BPE - #2234
Draft
ArthurZucker wants to merge 5 commits into
Draft
perf(tk-encode): SIMD-first encode pipeline — fused split cache + de-virtualized BPE#2234ArthurZucker wants to merge 5 commits into
ArthurZucker wants to merge 5 commits into
Conversation
…Split copy) classify_into_spans had the FSM write (start,end) into a thread-local Vec<Span> scratch, then a second pass copied scratch -> out as Split structs. Split is now repr(C) and layout-identical to atomsplit Span ((u32,u32)), so the FSM writes final Splits straight into out's spare capacity + set_len — one write, no scratch. HONEST NOTE: this is a cleanup, NOT a speedup. Re-profiling shows the split stage is FSM-bound (fsm_byte_level ~all of it), classify ~1.2%, and the copy was already vectorized/cheap. English split stays ~1.81 ns/B. An earlier claim that the copy was ~55% was a misread of the profiler's inclusive sample counts. The real split gap vs gigatoken (0.73) is the generic class-runs FSM vs their per-family hand-fused SIMD splitter.
…e probes register-only Classify->cache dataflow: the GPT byte-level split (classify_into_spans_keyed) packs each <=15-byte pre-token's bytes into a u128 (one wide load off the just-classified bytes) and carries it to the model. tokenize_pipeline derives the bucket hash with a hardware CRC and the FlatCache confirms the hit with a register 128-bit compare (s.key == key) — no ahash of the bytes, no kbytes memcmp, no key-arena load. Long (>15B) / non-GPT pre-tokens carry key 0 and fall back to the byte-verify path unchanged, so it's byte-exact (pack_key returns 0 for >15B to avoid prefix aliasing). Also: FSM writes spans straight into out (no scratch->Split copy; Split is now repr(C)==Span). Warm end-to-end (gpt2, ns/B, byte-exact vs reference): English 5.39->3.83 (1.41x), French 1.20x, Hindi 1.17x; cache+emit portion English 3.63->1.72 (2.1x). CJK ~flat (merge-bound; cache/emit is a small fraction there). Cost: split rises ~0.35 ns/B (the key-pack pass) — would vanish if packing were fused into the FSM span emission instead of a second loop.
…ed, no second loop) fsm_byte_level gains a const-generic KEYED path (fsm_byte_level_keyed) that writes pack_key(start,len) into a keys[] slice at the single span-emit site; cl100k/o200k pack from their emitted spans. classify_into_spans_keyed now hands the FSM the keys spare directly — no post-FSM re-walk. HONEST RESULT: byte-exact, all idsEq OK, but perf-NEUTRAL vs the separate loop — English split 2.11 -> 2.07 ns/B (was 1.75 pre-keying). The ~0.3 ns/B key cost is the per-pre-token 16-byte load + the keys[] write (~4MB/MB), which happen whether packing is in the FSM or a second loop; the FSM walks bytes via the tags array, not as 16-byte windows, so "bytes already in hand" doesn't remove the pack's own load. Eliminating it needs the fully-fused single-pass probe (no keys buffer) — the design we agreed isn't worth the generality cost. End-to-end warm unchanged (~1.44x English).
… split un-penalized The dispatch loop in encode_generic already holds the parent buffer (normalized_chunk), so it packs each pre-token's u128 key there (one wide load, safe mid-buffer) and passes it to tokenize_pipeline by value. No keys buffer materialized, and the split pays nothing. The cache then derives a CRC bucket + register 128-bit compare (no ahash, no memcmp, no key arena). Beats both earlier tries: carrying keys via a buffer (split rose to 2.07, buffer write) and packing at lookup (the pre-token slice can't wide-load safely -> slow memcpy fallback, o.mdlW ballooned). Packing where the parent buffer is hot gets the cheap wide load AND no buffer. Removes the now-dead keyed plumbing: fsm_byte_level_keyed, classify_into_spans_keyed, PreTokenizer::pre_tokenize_keyed (+ impls). gpt2 English warm o.tot 5.39 -> 3.63 (1.49x), split recovered to 1.76; byte-exact all langs.
The warm hot path called tokenize_pipeline once per pre-token through the PipelineModel enum, so
every pre-token paid an enum match, a thread-local PIPE_FLAT_CACHE borrow, and a retarget check.
PipelineBPE::tokenize_chunk now does those ONCE per chunk and runs a tight inner loop; the cache
exposes a small inlinable `get` (probe, returns arena offset/len) + `ids_tokens`, and the emit is a
repr(transparent) `&[u32]`->`&[PipelineToken]` memcpy (extend_from_slice), not a per-id map.
Back-to-back best-of-2, byte-exact all 11 langs x 5 tokenizers:
cl100k (llama3): -10..-17% (English 5.22->4.45, Hindi 3.79->3.13) - it emits more, smaller
pre-tokens and runs ignore_merges, so the per-pre-token overhead was a bigger share.
byte-level (gpt2): -2..-6%.
Dead end recorded in the code: fusing emit INTO the probe (get_into) made it too big to inline ->
a real call per pre-token -> +24% regression. Small `get` + memcpy in the loop body recovered it.
Inline-ids-in-slot was also measured and reverted: it regresses in our L2-resident regime (the wider
slot table costs more than the saved arena load, which is already hot here).
ArthurZucker
changed the base branch from
feat/train_encode_split
to
feat/bpe-cache
July 23, 2026 12:03
ArthurZucker
changed the base branch from
feat/bpe-cache
to
perf/tk-encode-pretok-base
July 23, 2026 12:09
SBrandeis
added a commit
that referenced
this pull request
Jul 29, 2026
* WordCache: one generation, evicting in place Replaces the 4-way set-associative table with an open-addressed one that makes room by overwriting rather than by rebuilding. Two things drove it, both measured in cache_bench on 3 models x 18 corpora: the 4-way table locks ~0.4% of words out on bucket conflicts that open addressing simply does not have, and the alternative that fixes that (PR #2234's FlatCache) can only reclaim arena space by compacting every live entry into a second set of buffers -- so it carries two of everything, 12.6 MB per instance against this one's 2.7 MB. Nothing here needs a second copy: - Eviction overwrites a slot in place. Slots go empty -> occupied and never back, so a probe walk is never cut short by a hole and there are no tombstones and no backward shift. The cost is that the walk needs a bound, hence WINDOW. - Slots are self-contained at 32 B, two to a cache line. A <= 15-byte word packs into the u128 key, so a hit is a register compare with no memcmp and no key arena; up to 3 ids ride in the slot, so most hits never touch the id arena either. - What still needs an arena takes a run from a free list per exact length (MAX_LENGTH bounds every run, so there is no rounding waste), and gives it back on eviction. The arenas therefore hold the live set instead of every insert, and never need compacting. Retention is FlatCache's idea made local: a full window evicts its coldest slot and halves the frequencies it passed, so an entry has to keep being used to keep its place. Measured against the 4-way it replaces, at the shipped 65,536 capacity: 0.000% steady-state miss rate against 0.413%, and x1.12 median warm replay across the 54 model x corpus pairs. Against FlatCache it is a near-tie on speed (x0.97) for 22% of the memory, and x11.6 ahead on the one corpus where FlatCache's arena budget forces it to cull. `get` now takes &mut self: it bumps the entry's frequency and unpacks inline ids into a scratch field to hand back a slice. The one call site already held a &mut binding, so it needed no change. * Apply suggestions from code review Co-authored-by: Simon Brandeis <33657802+SBrandeis@users.noreply.github.com> * document with AI * extract freq outside of the key * swap cache and ignore_merges * fast-track reject long keys * hash the key * no memcpy + Split type
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
1 participant
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.
Based on the pre-cache pipeline commit (
perf/tk-encode-pretok-base), so the diff is only this change. It is not based on #2223 (feat/bpe-cache): the two branches' pipelines diverged exactly here — this one is a 3-argFlatCachepath, #2223 is a 5-arg ScratchPool/word-cache path — so the change can't be a small diff on top of #2223 (it would be a reimplementation). Head-to-head perf of the full pipelines is in #2234's history / the comparison I posted.The change
memcmp), split pays nothing.retarget); now once per chunk, inner loop inlines pack→probe→emit. Emit is arepr(transparent)&[u32]→&[Token]memcpy.Byte-exact:
idsEq OKon 11 languages × 5 tokenizer families (gpt2, llama3/cl100k, deepseek, gpt-oss/o200k, glm).Perf (this change vs the same branch's pre-cache baseline, warm best-of, ns/byte)
The de-virtualization is the win, largest where per-pre-token overhead dominated:
Two measured dead-ends documented in the code so they aren't retried: inlining ids into the cache slot (regresses in our L2-resident regime) and fusing emit into the probe (un-inlinable → +24%).
Tests: 297 pass; 2 pre-existing
precompiled-normalizer failures unrelated to this change.