perf(bpe): wire the word cache into the batched model path, fused - #2308
perf(bpe): wire the word cache into the batched model path, fused#2308ArthurZucker wants to merge 3 commits into
Conversation
The GPT regexes are a scan with one unpredictable branch per token. `bitsplit` decides them from bitmaps instead: build a few `u64` masks over the bytes, then take token starts 64 bytes at a time with register ops rather than a branch per character. A new crate, wired in behind two match arms in `Split::pre_tokenize`. It only takes over where its bitstream build is SIMD -- `bitsplit::fast_builder()` -- because the portable builder is slower than the FSM it replaces, so every other target and every other regex keeps the existing atomsplit FSM. cl100k is gated on the standard digit cap of 3, the only variant with a bitstream splitter. `classify_into_spans_bits` is the `classify_into_spans` equivalent for a splitter that works off bitstreams: it hands the splitter two `u64` bitmaps alongside the tags. 367 tests pass, including the pre-tokenizer oracles that hold the FSM and the bitstream splitter to the same spans as the `regex` crate.
The pipeline has the entire span list before the model runs, so handing the spans over one at a time bought nothing and cost a virtual call, a slice, a `Result` and an output capacity check per pre-token -- on English, one round trip per ~5 bytes. `Model::tokenize_spans` takes them all. Its default is the loop it replaces, so a model only overrides it when it has per-chunk work to hoist; BPE does, and there the scratch destructuring and the output reservation move out of the loop. Slicing the chunk per pre-token also re-ran a UTF-8 boundary check on spans the pre-tokenizer had already cut on char boundaries. Ported onto #2241's `model.rs` rather than cherry-picked from #2304: that PR was written against the base before #2241 landed and restructured this file, so its patch no longer applies. Same change, current structure -- and the fold fast path it now sits beside came in with #2241. 367 tests pass.
#2306 gives the model a whole chunk of pre-tokens at a time by overriding `Model::tokenize_spans`. #2307 wires the word cache into `tokenize_pipeline`. Those are different functions: the pipeline calls the batched one, so landing both leaves the cache unreachable on the path that actually runs. This wires it into both. On the batched path the probe is fused into the emit. `WordCache::probe_emit` checks the home slot and, on an inline hit, stores all MAX_INLINE_IDS lanes straight at the caller's cursor, returning only the count. `lookup` hands back a `&[u32]` instead, which makes the caller re-read the slot to build a fat pointer and then copy a run whose length it learns at run time -- three trips over one 32-byte line that a single load already brought in. Anything that is not an inline home-slot hit falls back to the full window walk, reusing the placement so the word is never hashed twice. `lookup` and the walk are `#[inline]`. The output buffer is written through a cursor and closed with one `set_len`, so a hit costs no capacity check and no length store. It also reserves 2 ids per pre-token rather than 1: `spans.len()` is a *lower* bound (92% of English pre-tokens are one id, 98% at most two), so reserving it made the buffer grow -- and memcpy what it already held -- partway through most chunks. The fold stays in front of the cache: it answers a whole-vocabulary-entry word in one probe, which is cheaper than a cache probe, so folded words never enter the cache. They are already as cheap as a hit. Also drops `fold_by_flag`. Whether the config declared `ignore_merges` is a load -time question and does not belong in a per-pre-token branch: a declaring config asks every hit to fold, so every entry gets the bit; otherwise only the entries that prove they reduce to themselves earn it. `fold_id` is then one probe and one bit test with no policy left in it, and `ignore_merges` stops being a field. Fixes a sparse-id bug on the way: `prove_fold` bounded its walk by `vocab.len()`, the entry count, so any entry with an id above it was left unproven. Ids may be sparse, so the bound is the id space -- added as `BucketVocabStore::id_space`. No effect on gpt2, whose ids are dense, but the unified fold path above is wrong without it whenever a config leaves gaps. `PipelineToken` is asserted layout-identical to `u32`, since the fused probe writes ids through a pointer into the token buffer. 367 tests pass, 0 failures.
Benchmarkstokbench,
‡ ranges overlap 2306 — not resolvable at this noise floor (±5–20% per arm). Only the top three rows and gpt2/chinese have non-overlapping ranges. Geometric mean over the eight cells: 105.0 → 117.6 MB/s, 1.12×. Against 2279: 0.92×. What this recovers, and what it doesn'tRussian is fully closed — 1.29× / 1.37× over 2306, landing on 2279 (0.99× / 1.01×). Code moves 1.26× on gpt2 but stops well short of 2279 (0.73× / 0.66×). The residual on code is the fold and the cache both probing. On code a pretoken is usually not a whole vocabulary entry, so the fold probe misses and the cache probe follows it — two hashes of the same word where 2279, which has no fold at all for gpt2, pays one. This is the same shape #2302 already noted for a different gate: "the miss path then computes the gate twice." The fix is to share one hash between the two probes, which is a separate change and not attempted here. chinese is 0.94× on gpt2 and flat on llama-3: long CJK pretokens make the key expensive to build and repeat rarely, so the cache does not pay for itself there. A length gate on the cache probe is the obvious follow-up — 2279 caps at Exactness
Every cell above
Caveat on absolute numbersThese are tokbench's protocol: 10 kB chunks, ~1 MB of distinct text per cell, and the ids copied out into the harness buffer. They are not comparable to the ~1000 MB/s figures in #2279's description, which come from |
1 participant
Stacked on #2119 (
feat/train_encode_split), and meant to go on top of #2306 + #2307.The bug this exists to prevent
#2306 gives the model a whole chunk of pre-tokens at a time by overriding
Model::tokenize_spans. #2307 wires the word cache intotokenize_pipeline. Those are different functions, and the pipeline calls the batched one — so landing both as they stand leaves the cache unreachable on the path that actually runs. This wires it into both.That is the same failure mode the
eabddcdamerge on #2279 already produced once: work that lands in a function nothing calls.Fused probe
On the batched path the probe is fused into the emit.
WordCache::probe_emitchecks the home slot and, on an inline hit, stores allMAX_INLINE_IDSlanes straight at the caller's cursor and returns only the count.lookuphands back a&[u32], which makes the caller re-read the slot to build a fat pointer and then copy a run whose length it only learns at run time — three trips over one 32-byte line a single load already brought in. Anything that is not an inline home-slot hit falls back to the full window walk, reusing the placement so the word is never hashed twice.lookupand the walk are now#[inline].The output buffer is written through a cursor and closed with one
set_len, so a hit costs no capacity check and no length store. It also reserves 2 ids per pre-token instead of 1:spans.len()is a lower bound (92% of English pre-tokens are one id, 98% at most two), so reserving it made the buffer grow — and memcpy what it already held — partway through most chunks.Fold ordering
The fold stays in front of the cache. It answers a whole-vocabulary-entry word in one probe, cheaper than a cache probe, so folded words never enter the cache — they are already as cheap as a hit. This is also why the cache's marginal value looks small next to #2302: they are substitutes on the same population, not complements.
fold_by_flagremovedWhether the config declared
ignore_mergesis a load-time question and does not belong in a per-pre-token branch. A declaring config asks every hit to fold, so every entry gets the bit; otherwise only entries that prove they reduce to themselves earn it.fold_idbecomes one probe and one bit test with no policy in it, andignore_mergesstops being a field.Sparse-id fix
prove_foldbounded its walk byvocab.len()— the entry count — so any entry whose id exceeded it was left unproven. Ids may be sparse, so the bound is the id space, added asBucketVocabStore::id_space. No effect on gpt2 (dense ids), but the unified fold path above is wrong without it whenever a config leaves gaps; thebyte_level_ignore_merges_whole_wordtest (vocab id 300, 257 entries) catches exactly that.PipelineTokenis asserted layout-identical tou32, since the fused probe writes ids through a pointer into the token buffer.Status
cargo test -p tk-encode: 367 passed, 0 failed. Benchmarks posted below as a comment.