perf(bpe): Last missing nits to get up to speed - #2313
Open
ArthurZucker wants to merge 8 commits into
Open
Conversation
`tokenize_pipeline` hashed every fold-missing pretoken twice: `fold_id` hashed it for the vocabulary's MPHF, then `WordCache::lookup` hashed the same bytes again for its home slot and tag. `BucketVocabStore` and `WordCache` seed `ahash` with the same four constants, so the second pass recomputed a value the first had already produced. Share it: `hash_word` exposes the value, `get_bytes_foldable_hashed` and `lookup_hashed` take it, and the model computes it once per word. Verification is untouched. The vocabulary still compares the entry's bytes to the query in full, and the cache still compares its key, so ids cannot change -- only the duplicated hash goes away. A word over fifteen bytes still pays the cache's second, independently seeded discriminant hash, which is what makes its key 127 bits rather than 64. The sharing is only sound while both sides seed identically, so a test pins them together. Re-seeding either would leave the cache placing a word under one hash and looking it up under another: no wrong ids, but every lookup would miss and the cache would quietly stop working.
This was referenced Aug 7, 2026
Collaborator
Author
|
Rebased onto Benched the whole stack: geomean 1.1493x over 29 gpt2 cells, 29/29 faster, all ids identical, interleaved A/B/A/B/A/B. Full numbers and method on #2316. Note the stack has to land together to show that. |
A pretoken is short. English averages 4.83 bytes of it, code 4.08, and the `<|...|>` shapes in `added-special-dense` 2.29. Running aHash over that is most of what the fold probe costs, and it buys nothing: the vocabulary compares the entry's bytes anyway, so the hash only has to spread well enough for the MPHF to separate keys. For a word of seven bytes or fewer, pack the bytes and the length into a `u64` and mix them with one multiply. Seven, so the length still fits in the top byte, which is what keeps `"ab"` from colliding with `"ab\0"`. Longer words keep aHash, which mixes the length in itself. `word_hash` is now the one definition. `BucketVocabStore::build`, every probe, and the word cache's placement all go through it, so a pretoken probed in both tables is hashed once for the pair and the two cannot drift apart. The per-struct `RandomState` goes away with it: consistency came from carrying the hasher around, and now it comes from there being a single function. Verification is unchanged and stays exact -- the vocabulary still compares the entry's bytes to the query in full, the cache still compares its 128-bit key. Nothing verifies with `mix`, which is why it does not have to be a strong hash. Dropping the mixing altogether does not work: packed short keys share their high bytes and MPHF construction fails with "indistinguishable hashes in bucket". Note this is why `WordCache::lookup` has to call `placement_hash_of` rather than a hasher of its own: the `debug_assert` in `lookup_hashed` caught exactly that mistake while this was being written.
|
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. |
…2316) The fold probe was three dependent loads: the MPHF pilot, the entry, then the byte slab to compare the token against the query. The word cache does the same job in two, and the reason is layout, not luck -- its key lives in the slot it verifies, so nothing else has to be read. Give the vocabulary the same shape. `Entry` becomes `{ key, id }`, and `(start, len)` moves to a parallel `spans` array that only the reverse lookup and enumeration touch. A probe is now pilot + entry. Verification stays exact. A word of `INLINE_KEY_BYTES` or fewer has a key that *is* its bytes and its length, so comparing keys is proof of identity and the slab is never read. A longer word keys by aHash, which is not proof, so it still confirms against the slab -- the load it was paying anyway. So the saving lands exactly on the short pretokens that are the gap (english averages 4.83 bytes, code 4.08, `added-special-dense` 2.29) and nothing gives up the never-wrong guarantee. `LEN_TAG` now biases the length by one. A non-minimal MPHF returns padding slots, whose `Entry::default()` key is 0, and the probe rejects those with the same single compare it uses for everything else -- which only works while no real word can key to 0. The empty word keyed to exactly that before the bias. `key_and_hash` returns both halves so neither is recomputed: the model runs it once per word and hands the key and the hash to the fold probe and the hash to the cache.
`pipeline.rs` encodes through `tokenize_spans`, not `tokenize_pipeline`, and `tokenize_spans` was still hashing each word twice: `fold_id` for the vocabulary, then `cache.lookup` for the cache. Everything this PR does was landing only on `tokenize_pipeline`, which the encode loop does not call. Run `key_and_hash` once per word and hand the pair to `fold_id_keyed` and the hash to `lookup_hashed`, as `tokenize_pipeline` already does. `fold_id` had exactly one caller and folded into `fold_id_keyed` with it, taking a stale `#[allow(dead_code)]` with it. `the_batched_path_matches_the_reference` covers the path: ids compared against the legacy reference over thousands of spans in one chunk.
Collaborator
Author
Benched, with per-change attributionPushed the missing piece first ( Methodtokbench, gpt2, 29 verified cells, Rungs, all with the batched path wired:
Result (median ns/B over 4 interleaved runs)
Cross-checks against an independent earlier A/B of the whole stack: 1.1493x there vs 1.159x here. The split is mechanistic, not incidental
All ids identical to base in every cell of every round. Caveat on absolute numbersTaken with tokbench's pipeline adapter locally patched to |
`encode_generic` sizes a fresh `Vec` from the input length -- a guess -- and hands back a new allocation on every call. A caller encoding many inputs (a batch, a server loop, a benchmark) can reserve once and `clear()` between calls instead: fewer allocations, and no first-touch of the token array each time. Split it: `encode_generic_into` takes `&mut Vec<PipelineToken>`, and `encode_generic` becomes the allocating wrapper, so nothing existing changes. Measured on identical code with both forms available, tokbench gpt2, 29 cells against gigatoken: 0.9233x allocating vs 0.9536x reusing -- ~3% of geomean throughput, ~6% on english (3.188 -> 3.001 ns/B).
ArthurZucker
changed the title
The probe already stopped reading the byte slab for short words; this drops what remained.
`Entry` becomes `{ digest: u32, id: u32 }` -- 8 bytes, so twice as many per cache line -- and a
slot is verified by comparing 32 bits of the hash rather than the key, so nothing stores or
compares a length. `LookupKey` follows from `u128` to `u64`.
Together with what this branch already had -- one hash serving both the fold and the cache probe,
and the key living in the entry so the probe is two dependent loads instead of three -- the fold
path is now: one masked load and a multiply for the key, one pilot load, one entry load, one
compare.
Measured +6.1% geomean over tokbench's 29 gpt2 cells (median of four interleaved runs against
gigatoken inside each run). Biggest on the cells with the most pretokens per byte: chat-deepseek
1.25x, added-special-dense 1.16x, chat-chatml 1.16x, agentic-traces 1.16x, chat-llama3 1.16x.
Verification becomes probabilistic. A wrong id needs a 32-bit digest collision on a slot that is
occupied: (50257/65536) x 2^-32 = 1.8e-10 per DISTINCT pretoken -- per distinct, not per query,
since a word keys to the same slot and digest every time. Distinct pretokens grow ~1e4 per MB by
Heaps' law, so a 1 TB corpus reaches ~1e7-1e8 and expects 0.018 wrong ids; expecting one needs
5.5e9 distinct unseen byte strings, essentially the whole digest space. Exhaustively checking every
1-, 2- and 3-byte string (16,777,216 arbitrary-byte queries) gives 0 false positives with all
50,257 real tokens resolving.
`tokenize_spans` ran `fold_id_keyed` -- an MPHF probe -- ahead of the word cache, and reached
the cache through `lookup_keyed`, which walks the tag window, builds a `&[u32]` and `extend`s
it. So the expensive probe ran first and answered only the words that are their own vocabulary
entry, while every cache-servable word paid it for nothing. On a warm cache that is nearly all
of them.
Two changes, and they only pay together:
- the cache is reached through `probe_emit_keyed`, already in `word_cache`: a hit is one load
of the home slot and an unconditional store of its lanes, written straight at a running
cursor into the caller's buffer, so the ids never become a slice and the line is never read
twice. `output` is reserved once at two ids per span.
- the fold moves *behind* that probe, and a folded word is now inserted, so its second and
later occurrences come off the cache instead of re-probing the vocabulary.
The order is the whole point, and it follows whichever probe is cheaper. Measured on this
branch, against its own head:
fused emit, fold still first 1.013
fused emit, cache first 1.110 <- the reorder is +9.2% of that on its own
Reordering *without* the fused emit measures 0.951 -- a 5% regression -- because `lookup_keyed`
costs more than the fold probe the digest store made cheap. That is why the two land together.
ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against the branch
head with the LLC evicted between binaries. MB/s before -> after:
gpt2 english 786 -> 836 code 392 -> 446 dense 1572 -> 1651 chinese 900 -> 925
hindi 369 -> 507 thai 497 -> 610 korean 570 -> 592 russian 678 -> 702
greek 630 -> 658 arabic 604 -> 603
llama-3 english 794 -> 856 code 418 -> 432 dense 1442 -> 1621 chinese 957 -> 1001
hindi 652 -> 675 thai 803 -> 958 korean 667 -> 823 russian 790 -> 896
greek 776 -> 888 arabic 698 -> 886
warm geomean 1.110 (gpt2 1.095, llama-3 1.126), cold 1.058. Worst cell 0.999, best 1.374.
Byte-exact: token counts are unchanged on all 20 model x corpus pairs, and equal to c7ae7f4's
on the same corpora. `prove_fold` only sets the bit for an entry that merging its own text
reproduces, so a folded word and a merged word give the same ids -- the reorder moves which
path answers, not what it answers. 370 tests pass.
ArthurZucker
added a commit
that referenced
this pull request
Aug 7, 2026
`tokenize_spans` ran `fold_id_keyed` -- an MPHF probe, a pilot load plus a dependent entry load into the whole vocabulary -- ahead of `probe_emit_keyed`, which is one load of the home slot and an unconditional store of its lanes. The expensive probe went first and answered only the words that are their own vocabulary entry, while every word the cache was about to serve paid it for nothing. On a warm cache that is nearly all of them. The cache now goes first and the fold answers the miss, where it still beats running the merge engine. A folded word is inserted, so its second and later occurrences come off the cache instead of re-probing the vocabulary. The order follows whichever probe is cheaper, and here the fused emit already made that the cache. On the branch behind #2313, where `900b6a48`'s digest store makes the fold cheap and the cache is still reached through `lookup_keyed`, the same reordering measures 0.951 -- so it is the relative cost that decides, not the order itself. ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against this branch's head with the LLC evicted between binaries. MB/s before -> after: gpt2 english 1128 -> 1182 code 609 -> 611 dense 1534 -> 1621 chinese 882 -> 902 hindi 544 -> 595 thai 617 -> 654 korean 584 -> 610 russian 708 -> 767 greek 670 -> 710 arabic 621 -> 676 llama-3 english 1090 -> 1141 code 714 -> 730 dense 1412 -> 1490 chinese 914 -> 918 hindi 835 -> 810 thai 932 -> 1003 korean 798 -> 866 russian 898 -> 968 greek 870 -> 936 arabic 904 -> 984 warm geomean 1.053, cold 1.039. Against c7ae7f4 on the same box this takes the branch from 0.925 to 0.982. Reserving two ids per span instead of one, which c7ae7f4 does, measures +0.24% here -- inside the +-0.8% geomean noise floor -- so `975ed8df`'s one-per-span reservation stays. Byte-exact: token counts unchanged on all 20 model x corpus pairs and equal to c7ae7f4's. `prove_fold` only sets the bit for an entry that merging its own text reproduces, so a folded word and a merged word give the same ids; the reorder moves which path answers, not what it answers. 370 tests pass.
ArthurZucker
added a commit
that referenced
this pull request
Aug 7, 2026
* perf(bpe): hash a word once for both the fold and the cache probe `tokenize_pipeline` hashed every fold-missing pretoken twice: `fold_id` hashed it for the vocabulary's MPHF, then `WordCache::lookup` hashed the same bytes again for its home slot and tag. `BucketVocabStore` and `WordCache` seed `ahash` with the same four constants, so the second pass recomputed a value the first had already produced. Share it: `hash_word` exposes the value, `get_bytes_foldable_hashed` and `lookup_hashed` take it, and the model computes it once per word. Verification is untouched. The vocabulary still compares the entry's bytes to the query in full, and the cache still compares its key, so ids cannot change -- only the duplicated hash goes away. A word over fifteen bytes still pays the cache's second, independently seeded discriminant hash, which is what makes its key 127 bits rather than 64. The sharing is only sound while both sides seed identically, so a test pins them together. Re-seeding either would leave the cache placing a word under one hash and looking it up under another: no wrong ids, but every lookup would miss and the cache would quietly stop working. * fix(bpe): repair `tokenize_spans`, which does not compile (#2314) #2304 added `PipelineBPE::tokenize_spans` against the model as it stood then. #2241 replaced the merge engines and #2310 dropped `ignore_merges`, and because the two landed on separate branches the merge produced a `feat/train_encode_split` that does not build: error[E0425]: cannot find type `Span` in this scope error[E0026]: struct `BpeScratch` does not have fields named `merge_queue`, `skip`, `word` error[E0027]: pattern does not mention fields `symbols`, `queue` error[E0609]: no field `ignore_merges` on type `&PipelineBPE` error[E0061]: this method takes 3 arguments but 4 arguments were supplied Bring the batch loop back in line with `tokenize_pipeline`: destructure `{ symbols, queue, word_cache }`, run the fold, and call the current `merge_word(sequence, symbols, queue)` followed by `unmap`. The fold has to stay ahead of the cache probe, as it is in `tokenize_pipeline`. A word that is a foldable vocabulary entry is answered in one probe and never enters the cache; probing the cache first would fill it with words the fold already serves for free, and the two paths would disagree about its contents. `tokenize_spans` overrides a trait method whose default is the `tokenize_pipeline` loop, so the two can drift without anything failing to build -- that is how this got in. Add a test that runs thousands of spans through one chunk (repeats, so the cache fills and hits; folded words; merged words; punctuation runs; multi-byte scripts; a long unbroken run) and compares the ids to the legacy reference. * perf(bpe): pack a short word into its key instead of hashing it (#2315) A pretoken is short. English averages 4.83 bytes of it, code 4.08, and the `<|...|>` shapes in `added-special-dense` 2.29. Running aHash over that is most of what the fold probe costs, and it buys nothing: the vocabulary compares the entry's bytes anyway, so the hash only has to spread well enough for the MPHF to separate keys. For a word of seven bytes or fewer, pack the bytes and the length into a `u64` and mix them with one multiply. Seven, so the length still fits in the top byte, which is what keeps `"ab"` from colliding with `"ab\0"`. Longer words keep aHash, which mixes the length in itself. `word_hash` is now the one definition. `BucketVocabStore::build`, every probe, and the word cache's placement all go through it, so a pretoken probed in both tables is hashed once for the pair and the two cannot drift apart. The per-struct `RandomState` goes away with it: consistency came from carrying the hasher around, and now it comes from there being a single function. Verification is unchanged and stays exact -- the vocabulary still compares the entry's bytes to the query in full, the cache still compares its 128-bit key. Nothing verifies with `mix`, which is why it does not have to be a strong hash. Dropping the mixing altogether does not work: packed short keys share their high bytes and MPHF construction fails with "indistinguishable hashes in bucket". Note this is why `WordCache::lookup` has to call `placement_hash_of` rather than a hasher of its own: the `debug_assert` in `lookup_hashed` caught exactly that mistake while this was being written. * perf(bpe): put the key in the entry, so the fold probe is two loads (#2316) The fold probe was three dependent loads: the MPHF pilot, the entry, then the byte slab to compare the token against the query. The word cache does the same job in two, and the reason is layout, not luck -- its key lives in the slot it verifies, so nothing else has to be read. Give the vocabulary the same shape. `Entry` becomes `{ key, id }`, and `(start, len)` moves to a parallel `spans` array that only the reverse lookup and enumeration touch. A probe is now pilot + entry. Verification stays exact. A word of `INLINE_KEY_BYTES` or fewer has a key that *is* its bytes and its length, so comparing keys is proof of identity and the slab is never read. A longer word keys by aHash, which is not proof, so it still confirms against the slab -- the load it was paying anyway. So the saving lands exactly on the short pretokens that are the gap (english averages 4.83 bytes, code 4.08, `added-special-dense` 2.29) and nothing gives up the never-wrong guarantee. `LEN_TAG` now biases the length by one. A non-minimal MPHF returns padding slots, whose `Entry::default()` key is 0, and the probe rejects those with the same single compare it uses for everything else -- which only works while no real word can key to 0. The empty word keyed to exactly that before the bias. `key_and_hash` returns both halves so neither is recomputed: the model runs it once per word and hands the key and the hash to the fold probe and the hash to the cache. * perf(bpe): hash once on the batched path too `pipeline.rs` encodes through `tokenize_spans`, not `tokenize_pipeline`, and `tokenize_spans` was still hashing each word twice: `fold_id` for the vocabulary, then `cache.lookup` for the cache. Everything this PR does was landing only on `tokenize_pipeline`, which the encode loop does not call. Run `key_and_hash` once per word and hand the pair to `fold_id_keyed` and the hash to `lookup_hashed`, as `tokenize_pipeline` already does. `fold_id` had exactly one caller and folded into `fold_id_keyed` with it, taking a stale `#[allow(dead_code)]` with it. `the_batched_path_matches_the_reference` covers the path: ids compared against the legacy reference over thousands of spans in one chunk. * perf(pipeline): let the caller own the output buffer `encode_generic` sizes a fresh `Vec` from the input length -- a guess -- and hands back a new allocation on every call. A caller encoding many inputs (a batch, a server loop, a benchmark) can reserve once and `clear()` between calls instead: fewer allocations, and no first-touch of the token array each time. Split it: `encode_generic_into` takes `&mut Vec<PipelineToken>`, and `encode_generic` becomes the allocating wrapper, so nothing existing changes. Measured on identical code with both forms available, tokbench gpt2, 29 cells against gigatoken: 0.9233x allocating vs 0.9536x reusing -- ~3% of geomean throughput, ~6% on english (3.188 -> 3.001 ns/B). * merge: resolve duplicate tokenize_spans (keep the cache-aware, keyed version) The merge of #2313 into #2306 left two `tokenize_spans` definitions: git took both sides textually because they landed in different places. The stale one is #2306's, predating the word cache -- it destructures `BpeScratch { symbols, queue }` with no `word_cache` and calls the removed `fold_id`. Dropped it; kept the version that folds, probes the cache, and hashes each word once. * perf(bpe): carry pair ranks across multipass passes, and stop splicing Two changes to the multipass engine, ported from the target-encode work. **Ranks carried across passes.** A pass used to re-look-up every pair it walked over, so the passes summed to O(n^2) table lookups -- measured at **41.9 per merged word**. Only the pairs touching a merge's product actually change, so `ranks[i]` (the value of the pair `(symbols[i], symbols[i+1])`) is now seeded by `convert_multipass` -- which already looks every pair up, so seeding costs one store per pair and no extra lookup -- and carried. A pass copies the ranks it did not invalidate and pays `get_value` **twice per merge** instead of once per symbol. Finding the next target is then a scan of `ranks` with no lookups at all. `prods` holds the matching product ids, kept apart so the search array stays a dense `u32` of ranks alone. **No memmove per merge.** A merge used to splice: write the product, then `copy_within` symbols, ranks and products to close the gap -- three memmoves on every merge. Instead the word carries a `live` bitmap of which slots still hold a symbol and a merge clears one bit; "previous live" and "next live" are `leading_zeros`/`trailing_zeros`. The `MAX_MP = 24` bound is what makes this work: it puts the live set in one `u64`. Dead pair slots hold `u32::MAX`, which is also "does not merge", so the minimum search skips them for free. The superseded sweep machinery (`MergeState`, `merge_once`, `batch_merging_is_safe`, `NOT_LEGAL`) goes with it -- the batching those implemented is subsumed by carrying ranks. Byte-exact: `the_proven_fold_never_changes_the_ids` and `the_batched_path_matches_the_reference` both compare ids against the legacy reference. * perf(bpe): write cache hits straight at the output cursor A cache hit went through the tag row -- one load of 16 control bytes, a SIMD compare, then the slot -- and handed back a slice the caller walked. Both are avoidable for the common case, a word cached in its own home slot with at most three ids. `probe_emit_hashed` reads the home slot directly and stores all `MAX_INLINE_IDS` lanes unconditionally at a `*mut u32` the caller supplies, so the line is touched once and the ids never become a slice. Lanes past `ids_len` are dead: the caller advances its cursor by `ids_len` only, so the next word overwrites them or the final `set_len` cuts them off. A spilled or off-home slot falls back to the window walk without re-keying the word (`lookup_placed`, split out of `lookup_hashed`). `tokenize_spans` keeps a raw cursor and one capacity check per word covering both the fold's single write and the probe's lanes, and reserves `2 * spans.len() + MAX_INLINE_IDS` up front -- 92% of english pre-tokens are one id and 98% at most two, so the old `spans.len()` was a lower bound that made the buffer grow, and memcpy what it held, partway through most chunks. Deliberately NOT taken from the source branch: its `LookupKey` is a `u64`, which makes a hit on a word over seven bytes a 2^-64 proposition. This keeps the 128-bit key, so a hit stays exact for words up to fifteen bytes as before; only the emit is fused. Byte-exact: `the_batched_path_matches_the_reference` drives thousands of spans through this path, cache hits included, and compares ids to the legacy reference under debug assertions. * perf(bpe): take the target-encode word cache and vocabulary store u64 packed keys throughout, digest verification in the vocabulary store, and the pipelined probe helpers (probe_slot/entry_at/resolve_foldable). Accepts the exactness trades deliberately: a cache hit on a word over seven bytes is 2^-64, and an out-of-vocabulary pretoken can be mistaken for a vocabulary token at 2^-32, where both were previously impossible. * perf(bpe): reserve one id per span, and key from a masked load Two corrections to match the target-encode loop. `output.reserve(spans.len() + MAX_INLINE_IDS)`, not two apiece. Two was measured worse: the allocating entry point sizes its buffer at `len/4`, about one id per span, so asking for two forced a reallocation on every call that would not otherwise have happened. `key_and_hash_readable`: the span lies inside `chunk`, so everything up to the chunk's end is readable and a short word's key is one unaligned masked load instead of a head/tail stitch. Not done, and deliberately: wiring the pipelined probe (`probe_slot`/`entry_at`/ `resolve_foldable`). Those helpers exist but the source branch does not use them, having measured every version slower -- staging eight at a time 0.968x, carrying the next key a word early with a `prfm` 0.96x, carrying the probe answer a word early 0.95x, pairing two words 0.90x. The path is bound by instruction count, not latency. * Revert "perf(bpe): carry pair ranks across multipass passes, and stop splicing" Measured, and it does not pay: **+0.7% geomean** over tokbench's 29 gpt2 cells, inside the +-0.8% noise floor (median of four interleaved runs, five checksum-distinct binaries, ratio taken against gigatoken inside each run so it is immune to position drift). It is also lopsided rather than uniformly small: chat-llama3 1.15x, agentic-tools 1.14x, chat-deepseek 1.12x, code 1.05x, against added-normalized-dense 0.85x, hindi 0.94x, dense 0.96x. Taking table lookups from 41.9 per merged word to 2 per merge sounds decisive and is not, for an arithmetic reason: it only touches the pretokens that actually merge, which is ~8% on english. The other 92% never enter the engine, so the whole change is bounded by a small slice of the model phase. Not worth ~170 lines of engine rewrite plus two extra scratch buffers. The rest of the stack -- the fused probe, u64 keys, the digest store -- is unaffected and stays. * strip the rationale out of the code Comment blocks recording measured dead-ends, alternatives tried and their numbers belong in the PR, not in the source. Dropped every non-doc comment except `SAFETY` (load-bearing for the unsafe blocks), collapsed each doc block to its first line, kept doctests, and removed the batched-path test I had added. Net effect on the diff against this POC: +543/-227 before, +378/-442 now -- 64 lines fewer than the base rather than 300 more. * drop the port's own rationale comments, keep the ones that were already there The measured-dead-end essays and alternatives-tried notes this port added belong in the PR, not the source: 167 comment lines removed across word_cache, bucket_vocab_store and model. Every comment that existed in the base is preserved verbatim. The 33 base comment lines that no longer appear are the ones whose subject the port deleted -- `PLACEMENT_HASHER` and `DISCRIMINANT_HASHER` (gone with the u64 key), "the hasher is also stored on the struct" (the field is gone), `entries[slot] -> (offset, length, id)` (an entry is now `(digest, id)`), and `fold_id`'s doc (folded into `fold_id_keyed`). `SAFETY` comments and doctests are untouched. * wide 1.6.0 is yanked -> update cargo.toml to reference 1.5.0 * perf(bpe): probe the cache before the fold `tokenize_spans` ran `fold_id_keyed` -- an MPHF probe, a pilot load plus a dependent entry load into the whole vocabulary -- ahead of `probe_emit_keyed`, which is one load of the home slot and an unconditional store of its lanes. The expensive probe went first and answered only the words that are their own vocabulary entry, while every word the cache was about to serve paid it for nothing. On a warm cache that is nearly all of them. The cache now goes first and the fold answers the miss, where it still beats running the merge engine. A folded word is inserted, so its second and later occurrences come off the cache instead of re-probing the vocabulary. The order follows whichever probe is cheaper, and here the fused emit already made that the cache. On the branch behind #2313, where `900b6a48`'s digest store makes the fold cheap and the cache is still reached through `lookup_keyed`, the same reordering measures 0.951 -- so it is the relative cost that decides, not the order itself. ab_giga, 4 MB, single thread, warm, median of 10 rotated rounds interleaved against this branch's head with the LLC evicted between binaries. MB/s before -> after: gpt2 english 1128 -> 1182 code 609 -> 611 dense 1534 -> 1621 chinese 882 -> 902 hindi 544 -> 595 thai 617 -> 654 korean 584 -> 610 russian 708 -> 767 greek 670 -> 710 arabic 621 -> 676 llama-3 english 1090 -> 1141 code 714 -> 730 dense 1412 -> 1490 chinese 914 -> 918 hindi 835 -> 810 thai 932 -> 1003 korean 798 -> 866 russian 898 -> 968 greek 870 -> 936 arabic 904 -> 984 warm geomean 1.053, cold 1.039. Against c7ae7f4 on the same box this takes the branch from 0.925 to 0.982. Reserving two ids per span instead of one, which c7ae7f4 does, measures +0.24% here -- inside the +-0.8% geomean noise floor -- so `975ed8df`'s one-per-span reservation stays. Byte-exact: token counts unchanged on all 20 model x corpus pairs and equal to c7ae7f4's. `prove_fold` only sets the bit for an entry that merging its own text reproduces, so a folded word and a merged word give the same ids; the reorder moves which path answers, not what it answers. 370 tests pass. --------- Co-authored-by: Lysandre Debut <lysandre@huggingface.co>
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
3 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.
What
On the BPE pipeline path, every pretoken that misses the fold was hashed twice:
fold_id→BucketVocabStore::get_bytes_foldablehashes the word for the MPHF probe;WordCache::lookup→make_lookup_keyhashes the same bytes again for the home slot and tag.BucketVocabStore::SEEDSandword_cache::PLACEMENT_HASHERare seeded with the same fourconstants, so the second pass was recomputing a
u64the first had already produced.This shares that one value:
hash_wordexposes it,get_bytes_foldable_hashedandlookup_hashedaccept it, and
tokenize_pipelinecomputes it once per word.What this does not change
Verification is untouched, so ids cannot change. The vocabulary still compares the entry's bytes
to the query in full (perfect hashing only guarantees a valid slot for in-vocab keys), and the cache
still compares its 128-bit key. Only the duplicated hashing goes away.
A word longer than fifteen bytes still pays the cache's second, independently seeded
DISCRIMINANT_HASHERpass — that is what makes its key 127 bits rather than 64, and it isdeliberately left alone.
The public
lookup/get_bytes_foldableentry points keep their old signatures and hash the wordthemselves, so Unigram and WordPiece are unaffected.
The invariant, and why there is a test for it
Sharing the hash is only sound while both sides seed
ahashidentically. If someone re-seeds either,the failure is quiet and nasty: the cache would place a word under one hash and look it up under
another, producing no wrong ids but a permanent 0% hit rate — a large silent slowdown that no
correctness test would catch.
vocab_and_word_cache_hash_a_word_identicallypins the two together across empty, short, boundary(15/16 byte) and long words, so that change fails loudly here instead.
Verification
cargo test -p tk-encode --all-features: 357 passed. In particularthe_proven_fold_never_changes_the_idspasses, which compares pipeline ids against the legacyreference over eight strings covering folded words, merged words,
<|endoftext|>(whose entrydeliberately does not fold), CJK, code and a long repeated run.
because their fixtures are absent from a fresh worktree (
Os { code: 2, NotFound }):models::bpe::model::fold_tests::the_proven_fold_never_changes_the_ids(wants../data/gpt2.json— passes once fetched, as above) and
normalizers::precompiled::tests::pipeline_precompiled_matches_legacy.Neither is related to this change.
On throughput
No throughput number is claimed. The saving is one
ahashpass per fold-missing pretoken — around8% of pretokens on english, more on cache-heavy corpora — which is below the run-to-run noise floor
I measured on this machine (identical code, three runs: english 8.5% spread, math-latex 12.8%, code
9.2%, while chinese and hindi sit under 1.5%). The argument here is that the work is provably
redundant, not that a benchmark moved.
Scope
Deliberately minimal and independent of #2304 / #2308:
tokenize_spansor the fusedprobe_emitthat writes ids at the output cursor —those need the batched span path, which is perf(pipeline): give the model a whole chunk of pre-tokens at a time #2304's and does not exist on this base.
bytes. That is a separate change with a separate argument, since it would let an out-of-vocabulary
pretoken be mistaken for a vocabulary token at ~2^-32 where today that is impossible. Worth doing
on its own merits and its own review, not smuggled in behind "hash once".
PipelineTokenizer benchmark
10 / 10 models supported — PipelineTokenizer vs
tokenizersv0.23.1 (latest release) · whole corpus · ~10 kB chunks · cold caches · add_special_tokens on · 2/4-thread sweep per group900b6a480 · 2026-08-07 06:42 UTC· Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · 32 coresDecode — the release encodes each fixture's decode sample (
add_special_tokens=true) and both implementations decode those SAME ids withskip_special_tokens=false, so the comparison is decode alone. MB/s counts the input bytes the ids came from, the same denominator as the encode charts.vs base branch (
d64ad20b7) — per-model geomean ×speedup of this PR's PipelineTokenizer against the base branch's; regressions in red.Work vs base (
d64ad20b7, allocation lane):allocation lane measured on: Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · glibc 2.39
bert-base-uncased — normalizer-heavy WordPiece · ×3.50 vs v0.23.1 · ×0.71 vs base · decode ×0.61
Memory (RSS MB, load+encode+decode): v0.23.1 7+11+0 (peak 38) · Pipeline 4+4+0 (peak 42)
Allocations (encode pass, whole corpus): v0.23.1: 159.5M allocs, 52.04 GB allocated, peak live 37 MB; Pipeline: 61.6k allocs, 0.94 GB allocated, peak live 44 MB
Allocations (decode pass, decode sample): v0.23.1: 98.5M allocs, 1.26 GB allocated; Pipeline: 98.5M allocs, 1.26 GB allocated
deepseek-v4 — deepseek 3-regex split-heavy byte-level BPE · ×15.71 vs v0.23.1 · ×1.07 vs base · decode ×10.85
Memory (RSS MB, load+encode+decode): v0.23.1 57+0+0 (peak 75) · Pipeline 42+0+0 (peak 88)
Allocations (encode pass, whole corpus): v0.23.1: 153.2M allocs, 38.67 GB allocated, peak live 66 MB; Pipeline: 28.0k allocs, 0.23 GB allocated, peak live 74 MB
Allocations (decode pass, decode sample): v0.23.1: 15.3M allocs, 0.77 GB allocated; Pipeline: 7.4k allocs, 0.06 GB allocated
gemma-4 — byte-fallback BPE, Metaspace-style split (gemma-4) · ×2.94 vs v0.23.1 · ×0.96 vs base · decode ×1.79
Memory (RSS MB, load+encode+decode): v0.23.1 286+0+0 (peak 366) · Pipeline 1+0+0 (peak 365)
Allocations (encode pass, whole corpus): v0.23.1: 29.1M allocs, 23.69 GB allocated, peak live 310 MB; Pipeline: 65.2k allocs, 2.35 GB allocated, peak live 310 MB
Allocations (decode pass, decode sample): v0.23.1: 18.9M allocs, 1.79 GB allocated; Pipeline: 18.9M allocs, 1.79 GB allocated
gpt2 — gpt2 ByteLevel regex · ×27.90 vs v0.23.1 · ×1.06 vs base · decode ×33.35
Memory (RSS MB, load+encode+decode): v0.23.1 21+12+0 (peak 82) · Pipeline 6+4+0 (peak 84)
Allocations (encode pass, whole corpus): v0.23.1: 243.3M allocs, 38.17 GB allocated, peak live 79 MB; Pipeline: 34.6k allocs, 0.43 GB allocated, peak live 80 MB
Allocations (decode pass, decode sample): v0.23.1: 30.9M allocs, 1.36 GB allocated; Pipeline: 6.0k allocs, 0.07 GB allocated
gpt-oss — o200k-regex byte-level BPE (gpt-oss) · ×12.73 vs v0.23.1 · ×1.03 vs base · decode ×19.39
Memory (RSS MB, load+encode+decode): v0.23.1 236+0+0 (peak 310) · Pipeline 1+0+0 (peak 311)
Allocations (encode pass, whole corpus): v0.23.1: 127.8M allocs, 29.89 GB allocated, peak live 272 MB; Pipeline: 28.4k allocs, 0.23 GB allocated, peak live 272 MB
Allocations (decode pass, decode sample): v0.23.1: 14.5M allocs, 0.71 GB allocated; Pipeline: 7.4k allocs, 0.05 GB allocated
glm-5.2 — cl100k-variant regex byte-level BPE (glm-5.2) · ×16.10 vs v0.23.1 · ×1.03 vs base · decode ×16.04
Memory (RSS MB, load+encode+decode): v0.23.1 164+0+0 (peak 225) · Pipeline 1+0+0 (peak 226)
Allocations (encode pass, whole corpus): v0.23.1: 154.7M allocs, 33.00 GB allocated, peak live 212 MB; Pipeline: 29.5k allocs, 0.25 GB allocated, peak live 212 MB
Allocations (decode pass, decode sample): v0.23.1: 18.3M allocs, 0.90 GB allocated; Pipeline: 6.8k allocs, 0.05 GB allocated
llama-2 — model-bounded BPE, no pre-tokenizer · ×4.99 vs v0.23.1 · ×0.99 vs base · decode ×1.79
Memory (RSS MB, load+encode+decode): v0.23.1 20+6+0 (peak 63) · Pipeline 7+3+1 (peak 71)
Allocations (encode pass, whole corpus): v0.23.1: 63.0M allocs, 28.31 GB allocated, peak live 61 MB; Pipeline: 67.7k allocs, 0.56 GB allocated, peak live 72 MB
Allocations (decode pass, decode sample): v0.23.1: 35.7M allocs, 3.03 GB allocated; Pipeline: 35.7M allocs, 3.03 GB allocated
llama-3 — cl100k-regex byte-level BPE (llama-3), single regex · ×18.49 vs v0.23.1 · ×1.03 vs base · decode ×22.65
Memory (RSS MB, load+encode+decode): v0.23.1 68+2+0 (peak 90) · Pipeline 36+0+0 (peak 103)
Allocations (encode pass, whole corpus): v0.23.1: 152.3M allocs, 33.83 GB allocated, peak live 91 MB; Pipeline: 28.7k allocs, 0.24 GB allocated, peak live 92 MB
Allocations (decode pass, decode sample): v0.23.1: 17.3M allocs, 0.88 GB allocated; Pipeline: 9.6k allocs, 0.06 GB allocated
mistral-small-4 — tekken byte-level BPE, 1k added specials (mistral-small-4) · ×12.45 vs v0.23.1 · ×1.05 vs base · decode ×12.14
Memory (RSS MB, load+encode+decode): v0.23.1 151+0+0 (peak 189) · Pipeline 9+0+0 (peak 189)
Allocations (encode pass, whole corpus): v0.23.1: 135.2M allocs, 31.26 GB allocated, peak live 185 MB; Pipeline: 28.7k allocs, 0.24 GB allocated, peak live 185 MB
Allocations (decode pass, decode sample): v0.23.1: 16.6M allocs, 0.81 GB allocated; Pipeline: 7.1k allocs, 0.05 GB allocated
t5-base — Unigram + Metaspace · ×2.93 vs v0.23.1 · ×1.00 vs base · decode ×0.99
Memory (RSS MB, load+encode+decode): v0.23.1 31+10+0 (peak 56) · Pipeline 26+14+1 (peak 94)
Allocations (encode pass, whole corpus): v0.23.1: 260.7M allocs, 35.16 GB allocated, peak live 53 MB; Pipeline: 52.3M allocs, 4.77 GB allocated, peak live 107 MB
Allocations (decode pass, decode sample): v0.23.1: 12.0M allocs, 0.64 GB allocated; Pipeline: 12.0M allocs, 0.64 GB allocated