perf(bpe): WordCache - #2262
Open
SBrandeis wants to merge 14 commits into
Open
Conversation
|
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. |
SBrandeis
force-pushed
the
feat/bpe-cache-min
branch
from
July 30, 2026 11:09
4fc0e1e to
3b87e24
Compare
Plugs the WordCache into PipelineWordPiece the way BPE and Unigram already
use it: one lookup per pre-token, and on a miss the longest-match walk runs
and its ids go into the slot the lookup picked. The walk it skips is one
trie search per piece of the word, each over a fresh copy of what is left
to match, so a hit is worth a lot more than the lookup costs.
Capacity is DEFAULT_CACHE_CAPACITY and the cache is always on: WordPiece
has no cache knob in its config, unlike BPE and Unigram.
Measured on bert-base-uncased ("normalizer-heavy WordPiece") with
examples/fixture_bench.rs, in a worktree at HEAD so only this file differs
between the two binaries, alternating runs at process level (3 baseline, 2
cached, per-fixture medians over the 22 lang + modality fixtures):
median end-to-end +5.7% (base-vs-base noise floor: median 2.8%)
median model stage -32.1%
best: rus_Cyrl +22%, kor_Hang +17%, ell_Grek +15%, arb_Arab +14%
flat/slightly down: amh_Ethi -3.8%, agentic_swe -3.4%, cmn_Hani -1.4%
End-to-end gains stay modest because this model spends ~14 ns/byte in
normalization against 1-9 ns/byte in the model. ids_match held for all 22
fixtures on both cached runs, and encode RSS moved by ~0.1 MB.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
SBrandeis
commented
Jul 30, 2026
Comment on lines
+860
to
+861
| // The ids come back out of `output` because that is the only place both | ||
| // branches above leave them: `ignore_merges` never touches `word`. |
Contributor
Author
There was a problem hiding this comment.
Suggested change
| // The ids come back out of `output` because that is the only place both | |
| // branches above leave them: `ignore_merges` never touches `word`. |
SBrandeis
commented
Jul 30, 2026
| pub(crate) merge_queue: QuaternaryHeap<Merge>, | ||
| pub(crate) skip: Vec<Merge>, | ||
| pub(crate) word: Word, | ||
| /// Outlives the encode call that fills it, or it would never see a word twice. |
Contributor
Author
There was a problem hiding this comment.
Suggested change
| /// Outlives the encode call that fills it, or it would never see a word twice. |
SBrandeis
commented
Jul 30, 2026
|
|
||
| // The pool hands the SAME scratch to successive encodes. State left behind by one | ||
| // call — an undrained merge queue, a stale word buffer — would corrupt every call | ||
| // call (an undrained merge queue, a stale word buffer) would corrupt every call |
Contributor
Author
There was a problem hiding this comment.
Suggested change
| // call (an undrained merge queue, a stale word buffer) would corrupt every call | |
| // call — an undrained merge queue, a stale word buffer — would corrupt every call |
SBrandeis
marked this pull request as ready for review
July 30, 2026 17:07
SBrandeis
requested review from
ArthurZucker and
McPatate
and removed request for
ArthurZucker
July 30, 2026 17:11
SBrandeis
added a commit
that referenced
this pull request
Jul 31, 2026
A set-associative cache keyed on the word bytes, plus a scratch pile reused across encode calls, so BPE, WordPiece and Unigram all stop re-tokenizing words they have already seen. Squashed from feat/bpe-cache-min, which also carries #2261 (perf(scratch): reuse Scratch buffers across encode() calls).
ArthurZucker
reviewed
Aug 3, 2026
Comment on lines
+154
to
+157
| //! The walk reads tags, not entries: ruling a slot out costs one byte of memory | ||
| //! traffic rather than the 32 an entry takes. It rarely gets far. On running text a | ||
| //! word is usually in its home slot, or the slot after it, or not in the table at | ||
| //! all. |
Collaborator
There was a problem hiding this comment.
do we have a measure of this / a results? like usually when we run it ~is how many] probes?
Comment on lines
+162
to
+163
| //! compared. The key is the copy of the word kept inside the slot, and it is what | ||
| //! settles the question. |
Collaborator
There was a problem hiding this comment.
a reference? or an actual copy?
Comment on lines
+212
to
+220
| //! "hello" → [40,71,12,9,33] ┌────────────────┬────────────┬─────┐ | ||
| //! short word, too many ids │ h e l l o 5 │ ids are at │ ▒ │ | ||
| //! └────────────────┴────────────┴─────┘ | ||
| //! | ||
| //! "counterrevolutionary" → … ┌────────────────┬────────────┬─────┐ | ||
| //! word too long for the key │ hash of word H │ both are at│ ▒ │ | ||
| //! └────────────────┴────────────┴─────┘ | ||
| //! H says "a hash, where to look in | ||
| //! not the word" the buffers |
Collaborator
There was a problem hiding this comment.
this means its not super "friendly" with cache as you need at least 2 lookups, one in this table, one in the table that holds the ids!
Not something to fix, and it makes sense as of course long sequence of ids are bound to not fit.
But yeah, being able to do a single lookup can make a big dif!
Comment on lines
+230
to
+233
| //! The two overflow buffers ([`Arena`]) grow up to a budget and then stop. When an | ||
| //! entry is evicted, the space it used goes on a free list for its exact | ||
| //! length, so the next word of that shape reuses it. Nothing is ever compacted or | ||
| //! moved. Words longer than [`MAX_WORD_BYTES`] are not cached at all. |
Collaborator
There was a problem hiding this comment.
I think it can be interesting to compute the ratio / re-use ratio etc! like what is the distribution of length, etc
Comment on lines
+246
to
+247
| //! went. Two things let the plain rule get away with it: nothing is evicted until | ||
| //! all sixteen slots of a window are taken, and words repeat closely enough that a |
Collaborator
There was a problem hiding this comment.
how was 16 chosen btw? is it language specific?
ArthurZucker
reviewed
Aug 3, 2026
Collaborator
There was a problem hiding this comment.
Read this closely because the design is good and the doc is the only thing standing between it and the next person to touch it. Everything below is one coherent change plus a doc trim; the suggestions are meant to be applied as a batch, they don't compile one at a time.
The one thing to agree to first
A long word's key is built as (hash as u128) | KEY_IS_HASH — that uses 64 of the field's 128 bits and leaves 63 of them zero. Fill them, and the key comparison becomes the whole of confirming a hit. That single change deletes:
word_bytes_arena,word_off(),word_len()- the
hashedlocal and the&& …arena.get(…) == wordclause in the walk — one exit condition instead of two, and a hit costs the same for either kind of key PACKED_LEN_BITS,PACKED_LEN_MASKand itsconst _: () = assert!(and with them a silent-truncation edge:ids_len > 2047currently overflows into theword_lenfield)MAX_WORD_BYTESand the early return inlookup— with only a hash stored, word length stops matteringword: &'w [u8]fromPlacement, so it becomesCopyand both it andLookupshed a lifetime- the
!longclause inbuild_entry, so a long word with ≤3 ids now stores them inline — a straight win for exactly the CJK case the 3-id limit is weakest on - the
OptioninLookup::Miss, since no word is uncacheable any more
Cost: a second hash_one over the bytes, for words past fifteen bytes only, replacing a memcmp of those same bytes on every hit plus an arena write on every insert. I claimed that pays for itself; measured, it does so only up to a point. Warm lookup, ns per lookup, same fixed hash seeds on both sides so nothing but this change moves, min of 9 passes, identical hit rates:
| corpus | words | >15 B | old | new | |
|---|---|---|---|---|---|
| english | 1,095,696 | 0.5% | 18.11 | 17.28 | +4.6% |
| code | 23,533 | 16.2% | 5.70 | 5.59 | +2.0% |
| dense | 2,739 | 100% | 9.50 | 8.84 | +7.0% |
| russian | 14,209 | 34.4% | 6.23 | 5.85 | +6.1% |
| hindi | 14,310 | 31.7% | 6.45 | 6.05 | +6.3% |
| arabic | 19,693 | 5.6% | 5.93 | 5.83 | +1.7% |
| chinese | 723 | 99.9% | 22.43 | 24.92 | −11.1% |
By word length, holding everything else fixed: +24% at 16 B, +24% at 32 B, +21% at 64 B, +23% at 128 B, +8% at 192 B, +7% at 256 B — then it turns over: −18% at 384 B, −28% at 512 B, −58% at 1 KiB. Two full hash passes beat one pass plus a memcmp only while the bytes are cheap; past roughly 300 B the second pass costs more than the compare it removed. At 1 KiB that is two ~36 ns hash passes against one pass plus a ~9 ns memcmp, which is the shape measured.
Every corpus above has ~0% of its pre-tokens past that crossover except Chinese, which has 99.9% of them past it — hence the one negative row. So: a clear win wherever pre-tokens are short, a real loss on long CJK runs, which is the same regime the three-inline-ids limit is already weakest on.
Two honest caveats. The harness splits on whitespace, so the Chinese magnitude is indicative rather than what the real pre-tokenizer would show — long CJK pre-tokens are genuinely a thing, but not necessarily 700 bytes of them. And the fix if that regime matters is a single-pass 128-bit hash (xxh3-128), which takes the tail back at the cost of a dependency; I would not add it on this evidence alone.
What it buys is not exactness — it's near exactness. Two words whose 127 bits agree would trade ids. That rate is below the one at which the machine corrupts the ids in DRAM on the way past, but it is not zero, and the header now says so out loud instead of implying "never". That's the call to make; everything else follows from it.
The free lists
free: Box<[Vec<u32>]> is 1025 empty Vecs per arena — 2050 allocations to hold a free list. The cache is allowed to forget, so: bump-allocate, and when the budget is spent clear the tag row and start over. That deletes release, reclaim, the per-length filing, and the aliasing invariant that a_hit_never_returns_another_words_ids was written to defend. A flush is one memset of n_slots bytes.
The doc
~300 lines for ~440 of code, and a quarter of it explains a use-counter eviction policy that isn't in the file. Trimmed to ~95: one diagram showing that a single hash yields the tag and the index and that both rows take the same index, then the decisions that aren't recoverable from the code — why the index takes the low bits and the tag the top eight (and why that costs nothing), why EMPTY is a value and not a flag bit, why the key is 128 bits wide, and that three inline ids is what 32 bytes left over rather than a tuned number. Named the technique properly too: open addressing, linear probing, bounded probe window, SwissTable's H1/H2 split.
Two things I did not touch
- The two-phase
lookup/insertAPI stays. It's arguable — a miss is the path that then runsmerge_all, so hashing twice would be noise on it — but it isn't wrong, and it's not what this review is about. - The
memcpyTODO inpack_wordstays verbatim. It's the best remaining lead in the file: a third of key construction, and it goes away if the caller hands over the sixteen bytes starting at the word.
Applied all of this locally and ran it — 11/11 tests pass, including the two rewritten ones. The three caller suggestions are what keeps the batch compiling.
Comment on lines
+1
to
+301
| //! A table that remembers which token ids a word encodes to, so a model only has | ||
| //! to work it out once. | ||
| //! | ||
| //! # Context | ||
| //! | ||
| //! A tokenizer turns text into *token ids*: small integers a language model reads. | ||
| //! It gets there in three stages: | ||
| //! - It **normalizes** the text (for example: lowercasing), | ||
| //! - It **pre-tokenizes** the normalized text (ie, cuts it into small pieces, usually words or fragments of words), | ||
| //! - It runs a tokenization **model** over each piece to produce token ids | ||
| //! | ||
| //! This module calls those pieces *words*, because that is what they nearly always are. | ||
| //! | ||
| //! ```text | ||
| //! "The cat sat on The Mat" | ||
| //! │ | ||
| //! │ normalize, then pre-tokenize | ||
| //! ▼ | ||
| //! "the" "cat" "sat" "on" "the" "mat" ← words | ||
| //! │ │ │ │ │ │ | ||
| //! ▼ ▼ ▼ ▼ ▼ ▼ | ||
| //! ┌─────────────────────────────────────────────┐ | ||
| //! │ Tokenization Model │ | ||
| //! └─────────────────────────────────────────────┘ | ||
| //! │ │ │ │ │ │ | ||
| //! ▼ ▼ ▼ ▼ ▼ ▼ | ||
| //! [12] [87] [43] [9] [12] [64] ← token ids, made up here | ||
| //! ``` | ||
| //! | ||
| //! The model stage is expensive, due to the underlying algorithms: | ||
| //! | ||
| //! - [**BPE**](crate::models::bpe::BPE) (byte-pair encoding) starts from the word's individual | ||
| //! bytes and repeatedly glues the best-ranked neighboring pair together, until no pair | ||
| //! left in the word can be merged. Every merge changes the word and the search | ||
| //! for the next best pair starts again ([`Word::merge_all`](crate::models::bpe::Word::merge_all)). | ||
| //! - [**Unigram**](crate::models::unigram::Unigram) considers the many ways of cutting the word | ||
| //! into pieces the vocabulary knows, scores them, and keeps the best. That is a search over a | ||
| //! lattice of candidate cuts ([`Lattice::viterbi`](crate::models::unigram::Lattice::viterbi)). | ||
| //! - [**WordPiece**](crate::models::wordpiece::WordPiece) walks the word from the front, taking | ||
| //! the longest piece the vocabulary holds, then carries on from where that piece ended. That is | ||
| //! one search per piece of the word. | ||
| //! | ||
| //! Two things make that expense avoidable: | ||
| //! - Words repeat, in every kind of text. | ||
| //! - A word always encodes to the same ids. | ||
| //! | ||
| //! So a word only ever has to go through the model once. Cache the ids that came out, and | ||
| //! every later occurrence of that word skips the model altogether. | ||
| //! | ||
| //! This module is the table that does the remembering, in bounded memory: | ||
| //! | ||
| //! ```text | ||
| //! "The cat sat on The Mat" | ||
| //! │ | ||
| //! │ normalize, then pre-tokenize | ||
| //! ▼ | ||
| //! "the" "cat" "sat" "on" "the" "mat" ← words, one at a time | ||
| //! │ │ │ │ │ │ | ||
| //! ▼ ▼ ▼ ▼ ▼ ▼ | ||
| //! ┌─────────────────────────────────────────────┐ | ||
| //! │ WordCache: have I encoded this word yet? │ | ||
| //! └─────────────────────────────────────────────┘ | ||
| //! │ │ | ||
| //! miss │ │ hit | ||
| //! ▼ │ | ||
| //! ┌─────────────────────┐ │ | ||
| //! │ Model (expensive) │ │ | ||
| //! └─────────────────────┘ │ | ||
| //! │ │ | ||
| //! │ store the ids │ the ids, straight from the table | ||
| //! └─────┬──────────┘ | ||
| //! ┌──────┬──────┼──────┬──────┬──────┐ | ||
| //! ▼ ▼ ▼ ▼ ▼ ▼ | ||
| //! [12] [87] [43] [9] [12] [64] ← the same ids, most of them for free | ||
| //! ▲ | ||
| //! └ the second "the": a cache hit | ||
| //! ``` | ||
| //! | ||
| //! # What the cache is allowed to get wrong | ||
| //! | ||
| //! The cache is free to *forget*: if a word's ids are gone, the model works them | ||
| //! out again. What it must never do is hand back the *wrong* ids. Every | ||
| //! trade-off below spends the first freedom and none of the second. | ||
| //! | ||
| //! # The table | ||
| //! | ||
| //! The cache is one long row of numbered **slots**. Each slot holds one word and | ||
| //! the ids that word encodes to. The number of slots is fixed when the cache is | ||
| //! built and never changes, which is what keeps the memory bounded. | ||
| //! | ||
| //! To decide where a word belongs, the cache turns the word into a number, its | ||
| //! **hash**. Two different pieces of that number are used, and they never overlap: | ||
| //! | ||
| //! ```text | ||
| //! "the" | ||
| //! │ | ||
| //! │ hash | ||
| //! ▼ | ||
| //! ┌─────────────────────────────────────────┐ | ||
| //! │ 1010 0111 ................... 0000 0101 │ | ||
| //! └────┬───────────────────────────────┬────┘ | ||
| //! │ the top 8 bits │ the bottom bits | ||
| //! ▼ ▼ | ||
| //! the tag: A7 the home slot: 5 | ||
| //! ``` | ||
| //! | ||
| //! The **home slot** is where the word would like to live. The **tag** is a | ||
| //! one-byte summary of the word. Tags are kept in their own row, one byte per | ||
| //! slot, beside the row of entries: | ||
| //! | ||
| //! ```text | ||
| //! slot index: 0 1 2 3 4 5 6 7 ... | ||
| //! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ | ||
| //! tags │ · │ 9B │ · │ · │ C4 │ A7 │ 31 │ A7 │ 1 byte each | ||
| //! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ | ||
| //! entries │ │"of" │ │ │"cat"│"the"│"sat"│"hat"│ 32 bytes each | ||
| //! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ | ||
| //! ▲ | ||
| //! home slot of "the", | ||
| //! and of "hat" as well | ||
| //! | ||
| //! · = empty slot | ||
| //! ``` | ||
| //! | ||
| //! A word does not always get its home slot: another word may have taken it | ||
| //! first, which is what happened to "hat" above. So the word may be in its home | ||
| //! slot or in any of the fifteen slots after it. Those sixteen slots are the | ||
| //! word's **window**, and the word is either in there or nowhere | ||
| //! ([`WALK_WINDOW`]). | ||
| //! | ||
| //! # Looking a word up | ||
| //! | ||
| //! [`WordCache::lookup`] walks the window one tag at a time, starting at the home | ||
| //! slot, and answers both of a lookup's questions on the way: which slot holds | ||
| //! this word, and where would this word go if none of them does. | ||
| //! | ||
| //! ```text | ||
| //! looking up "hat" (tag A7, home slot 5) | ||
| //! | ||
| //! slot: 5 6 7 8 9 10 11 12 ... | ||
| //! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ | ||
| //! tags │ A7 │ 31 │ A7 │ · │ 5F │ A7 │ · │ · │ | ||
| //! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ | ||
| //! entries │"the"│"sat"│"hat"│ │"on" │"mat"│ │ │ | ||
| //! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ | ||
| //! ▲ ▲ ▲ ▲ | ||
| //! │ │ │ └ tagged A7 too, but past the | ||
| //! │ │ │ empty slot: out of reach | ||
| //! │ │ └ empty: the search stops here | ||
| //! │ └ tag matches, and the key says "hat": found it | ||
| //! └ tag matches, but the key says "the", so keep looking | ||
| //! ``` | ||
| //! | ||
| //! The walk reads tags, not entries: ruling a slot out costs one byte of memory | ||
| //! traffic rather than the 32 an entry takes. It rarely gets far. On running text a | ||
| //! word is usually in its home slot, or the slot after it, or not in the table at | ||
| //! all. | ||
| //! | ||
| //! A matching tag is a hint, not an answer. A tag is only one byte of the hash, so | ||
| //! about one slot in 255 matches a word that is not there, which is exactly what | ||
| //! slot 5 does above. Each matching slot is therefore read and its **key** | ||
| //! compared. The key is the copy of the word kept inside the slot, and it is what | ||
| //! settles the question. | ||
| //! | ||
| //! The search stops at the first empty slot. That is safe because storing stops | ||
| //! there too: a word is always put in the first empty slot of its window, so no | ||
| //! word can ever sit past one. This is why slot 10 is out of reach: whatever is | ||
| //! in it, it is not "hat", because "hat" would have taken slot 8. | ||
| //! | ||
| //! A lookup ends one of two ways ([`Lookup`]): | ||
| //! | ||
| //! - **Hit**: the ids, ready to use. Nothing is written: a hit leaves the table | ||
| //! exactly as it was. | ||
| //! - **Miss**: a [`Placement`], the slot this word should take. The caller runs | ||
| //! the model, then hands the placement back to [`WordCache::insert`] along with | ||
| //! the ids. The placement carries the hash work the lookup already did, so | ||
| //! storing the word costs no second hash and no second read of the tags. | ||
| //! | ||
| //! # Storing a word | ||
| //! | ||
| //! [`WordCache::insert`] takes the placement from the miss and the ids the model | ||
| //! produced, and does four things: | ||
| //! | ||
| //! 1. **Builds the entry.** The word and its ids go in the slot if they fit; what | ||
| //! does not fit is copied into an overflow buffer (below). If those buffers are | ||
| //! full, the insert is dropped and nothing changes, since the cache is allowed | ||
| //! to forget. | ||
| //! 2. **Clears the old entry, if there was one.** When the window had no empty | ||
| //! slot the placement points at the home slot, and that word's ids are lost. | ||
| //! Any overflow space it was using is handed back. | ||
| //! 3. **Writes the tag**, which is what makes the slot findable. | ||
| //! 4. **Writes the entry.** | ||
| //! | ||
| //! # What a slot holds | ||
| //! | ||
| //! A slot is 32 bytes: sixteen for the key, twelve for the ids, one that says how | ||
| //! many ids there are, and three of padding. It comes in three shapes | ||
| //! ([`CachedWord`]). The first is the common case, and the reason the sizes are what | ||
| //! they are: the whole answer (word and ids) is right there, so a hit reads one | ||
| //! slot and stops. | ||
| //! | ||
| //! ```text | ||
| //! ├───── key ─────┤├─── ids ───┤├─────┤ | ||
| //! how many | ||
| //! | ||
| //! "the" → [464] ┌────────────────┬────────────┬─────┐ | ||
| //! short word, up to 3 ids │ t h e 3 │ 464 │ 1 │ | ||
| //! └────────────────┴────────────┴─────┘ | ||
| //! the word, then the ids | ||
| //! its length themselves | ||
| //! | ||
| //! "hello" → [40,71,12,9,33] ┌────────────────┬────────────┬─────┐ | ||
| //! short word, too many ids │ h e l l o 5 │ ids are at │ ▒ │ | ||
| //! └────────────────┴────────────┴─────┘ | ||
| //! | ||
| //! "counterrevolutionary" → … ┌────────────────┬────────────┬─────┐ | ||
| //! word too long for the key │ hash of word H │ both are at│ ▒ │ | ||
| //! └────────────────┴────────────┴─────┘ | ||
| //! H says "a hash, where to look in | ||
| //! not the word" the buffers | ||
| //! | ||
| //! ▒ = "not in the slot, follow the offsets" | ||
| //! ``` | ||
| //! | ||
| //! A word of up to 15 bytes is kept whole, with its length in the sixteenth byte. | ||
| //! A longer word cannot fit, so the slot keeps the word's hash instead and the | ||
| //! bytes go to a buffer. Two different long words can hash to the same number, so | ||
| //! for those, and only those, a match is confirmed by comparing the bytes. | ||
| //! | ||
| //! The two overflow buffers ([`Arena`]) grow up to a budget and then stop. When an | ||
| //! entry is evicted, the space it used goes on a free list for its exact | ||
| //! length, so the next word of that shape reuses it. Nothing is ever compacted or | ||
| //! moved. Words longer than [`MAX_WORD_BYTES`] are not cached at all. | ||
| //! | ||
| //! # Which entry gets evicted | ||
| //! | ||
| //! When a word's window is full, one of its entries gets evicted, and it is whatever | ||
| //! sits in the home slot. Nothing is measured and nothing is ranked: the newcomer | ||
| //! takes the slot it wanted in the first place. | ||
| //! | ||
| //! A cleverer rule is possible, and this module used to have one: a use counter per | ||
| //! entry, the least-used slot of the window loses it, and the counters fade so that | ||
| //! a word that was busy early in a document cannot keep its slot forever. Measured | ||
| //! against the plain rule, it did keep slightly more of the words worth keeping, and | ||
| //! it lost more than it gained to the counter that every hit had to write, so it | ||
| //! went. Two things let the plain rule get away with it: nothing is evicted until | ||
| //! all sixteen slots of a window are taken, and words repeat closely enough that a | ||
| //! word's next use usually comes round before anything has had the chance to evict | ||
| //! it. | ||
| //! | ||
| //! # Why it is built this way | ||
| //! | ||
| //! - **Tags live in their own row** so that a walk over sixteen slots reads sixteen | ||
| //! bytes, which is one or two cache lines. The entries are 32 bytes each, so the | ||
| //! same sixteen slots would be 512 bytes and eight cache lines. | ||
| //! - **A tag is one byte of the hash**, so a slot that cannot hold the word is ruled | ||
| //! out without reading it. All eight bits carry hash: emptiness needs a value no | ||
| //! live tag can take rather than a bit of its own, and one value is cheaper than | ||
| //! one bit ([`make_tag`]). | ||
| //! - **The tag uses the top bits of the hash** because the bottom bits already | ||
| //! chose the home slot, and every slot in a window would share those. | ||
| //! - **A window is sixteen slots**: how far a word may end up from its home slot | ||
| //! before the cache stops looking and evicts instead. Long enough that a full | ||
| //! window is uncommon, short enough that walking a full one stays cheap. | ||
| //! - **The key is 128 bits** because 15 bytes of word plus one byte of length | ||
| //! covers nearly every word whole, and comparing it is then one comparison of one | ||
| //! value instead of following a pointer to bytes elsewhere in memory. | ||
| //! - **The length is part of the key**, or `"a"` and `"a\0"` would look identical. | ||
| //! - **Up to three ids fit in the slot** because most words encode to one, two or | ||
| //! three, and that makes a hit a single read. | ||
| //! - **The number of slots is a power of two**, so picking the home slot is one | ||
| //! bit operation rather than a division. So is folding a step of the walk back | ||
| //! into the table when a window runs off the end. | ||
| //! - **A miss hands back a placement** because the lookup already knows the hash, | ||
| //! the tag and the slot to use, and making the insert work them out again would | ||
| //! double that cost on every new word. | ||
| //! - **Freed overflow space is filed by exact length**: every run is at most | ||
| //! [`MAX_WORD_BYTES`] long, so there can be one free list per length and a freed | ||
| //! run always fits the next word of that shape exactly. | ||
| //! | ||
| //! # Where the ideas come from | ||
| //! | ||
| //! - [Swiss Tables] is where the control byte comes from: one byte per slot, holding | ||
| //! a flag bit and seven bits of hash. Here the whole byte is hash, because the flag | ||
| //! bit was there to be tested sixteen lanes at a time and this walk tests one. | ||
| //! - [gigatoken] is a BPE tokenizer with a pre-token cache built from the same | ||
| //! parts: `u128` packed keys with the length in the top byte, self-contained | ||
| //! 32-byte entries, ids inline. It never evicts (it doubles at 3/4 load) and | ||
| //! leans on huge pages and prefetching, because its table is sized for DRAM | ||
| //! rather than for a CPU cache. | ||
| //! - [TinyLFU] (Einziger, Friedman & Manes) is the use-counter-and-fade rule this | ||
| //! module tried for eviction and dropped, as its reference implementation | ||
| //! [Caffeine] does it. Kept here for whoever wants to try it again. | ||
| //! - [huggingface/tokenizers#2234] is an open-addressed cache for this same encode | ||
| //! pipeline, arrived at in parallel, fused into the pre-tokenizer's split loop. | ||
| //! | ||
| //! [Swiss Tables]: https://abseil.io/about/design/swisstables | ||
| //! [gigatoken]: https://github.com/marcelroed/gigatoken | ||
| //! [TinyLFU]: https://arxiv.org/abs/1512.00727 | ||
| //! [Caffeine]: https://github.com/ben-manes/caffeine/blob/master/caffeine/src/main/java/com/github/benmanes/caffeine/cache/FrequencySketch.java | ||
| //! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 |
Collaborator
There was a problem hiding this comment.
Header: ~300 lines -> ~95. Keeps the one diagram, the non-obvious decisions, the prior art. Drops the pipeline explainer, the walk/slot diagrams, the eviction-policy essay (it describes a policy that isn't in the file) and the TinyLFU citation that goes with it.
Suggested change
| //! A table that remembers which token ids a word encodes to, so a model only has | |
| //! to work it out once. | |
| //! | |
| //! # Context | |
| //! | |
| //! A tokenizer turns text into *token ids*: small integers a language model reads. | |
| //! It gets there in three stages: | |
| //! - It **normalizes** the text (for example: lowercasing), | |
| //! - It **pre-tokenizes** the normalized text (ie, cuts it into small pieces, usually words or fragments of words), | |
| //! - It runs a tokenization **model** over each piece to produce token ids | |
| //! | |
| //! This module calls those pieces *words*, because that is what they nearly always are. | |
| //! | |
| //! ```text | |
| //! "The cat sat on The Mat" | |
| //! │ | |
| //! │ normalize, then pre-tokenize | |
| //! ▼ | |
| //! "the" "cat" "sat" "on" "the" "mat" ← words | |
| //! │ │ │ │ │ │ | |
| //! ▼ ▼ ▼ ▼ ▼ ▼ | |
| //! ┌─────────────────────────────────────────────┐ | |
| //! │ Tokenization Model │ | |
| //! └─────────────────────────────────────────────┘ | |
| //! │ │ │ │ │ │ | |
| //! ▼ ▼ ▼ ▼ ▼ ▼ | |
| //! [12] [87] [43] [9] [12] [64] ← token ids, made up here | |
| //! ``` | |
| //! | |
| //! The model stage is expensive, due to the underlying algorithms: | |
| //! | |
| //! - [**BPE**](crate::models::bpe::BPE) (byte-pair encoding) starts from the word's individual | |
| //! bytes and repeatedly glues the best-ranked neighboring pair together, until no pair | |
| //! left in the word can be merged. Every merge changes the word and the search | |
| //! for the next best pair starts again ([`Word::merge_all`](crate::models::bpe::Word::merge_all)). | |
| //! - [**Unigram**](crate::models::unigram::Unigram) considers the many ways of cutting the word | |
| //! into pieces the vocabulary knows, scores them, and keeps the best. That is a search over a | |
| //! lattice of candidate cuts ([`Lattice::viterbi`](crate::models::unigram::Lattice::viterbi)). | |
| //! - [**WordPiece**](crate::models::wordpiece::WordPiece) walks the word from the front, taking | |
| //! the longest piece the vocabulary holds, then carries on from where that piece ended. That is | |
| //! one search per piece of the word. | |
| //! | |
| //! Two things make that expense avoidable: | |
| //! - Words repeat, in every kind of text. | |
| //! - A word always encodes to the same ids. | |
| //! | |
| //! So a word only ever has to go through the model once. Cache the ids that came out, and | |
| //! every later occurrence of that word skips the model altogether. | |
| //! | |
| //! This module is the table that does the remembering, in bounded memory: | |
| //! | |
| //! ```text | |
| //! "The cat sat on The Mat" | |
| //! │ | |
| //! │ normalize, then pre-tokenize | |
| //! ▼ | |
| //! "the" "cat" "sat" "on" "the" "mat" ← words, one at a time | |
| //! │ │ │ │ │ │ | |
| //! ▼ ▼ ▼ ▼ ▼ ▼ | |
| //! ┌─────────────────────────────────────────────┐ | |
| //! │ WordCache: have I encoded this word yet? │ | |
| //! └─────────────────────────────────────────────┘ | |
| //! │ │ | |
| //! miss │ │ hit | |
| //! ▼ │ | |
| //! ┌─────────────────────┐ │ | |
| //! │ Model (expensive) │ │ | |
| //! └─────────────────────┘ │ | |
| //! │ │ | |
| //! │ store the ids │ the ids, straight from the table | |
| //! └─────┬──────────┘ | |
| //! ┌──────┬──────┼──────┬──────┬──────┐ | |
| //! ▼ ▼ ▼ ▼ ▼ ▼ | |
| //! [12] [87] [43] [9] [12] [64] ← the same ids, most of them for free | |
| //! ▲ | |
| //! └ the second "the": a cache hit | |
| //! ``` | |
| //! | |
| //! # What the cache is allowed to get wrong | |
| //! | |
| //! The cache is free to *forget*: if a word's ids are gone, the model works them | |
| //! out again. What it must never do is hand back the *wrong* ids. Every | |
| //! trade-off below spends the first freedom and none of the second. | |
| //! | |
| //! # The table | |
| //! | |
| //! The cache is one long row of numbered **slots**. Each slot holds one word and | |
| //! the ids that word encodes to. The number of slots is fixed when the cache is | |
| //! built and never changes, which is what keeps the memory bounded. | |
| //! | |
| //! To decide where a word belongs, the cache turns the word into a number, its | |
| //! **hash**. Two different pieces of that number are used, and they never overlap: | |
| //! | |
| //! ```text | |
| //! "the" | |
| //! │ | |
| //! │ hash | |
| //! ▼ | |
| //! ┌─────────────────────────────────────────┐ | |
| //! │ 1010 0111 ................... 0000 0101 │ | |
| //! └────┬───────────────────────────────┬────┘ | |
| //! │ the top 8 bits │ the bottom bits | |
| //! ▼ ▼ | |
| //! the tag: A7 the home slot: 5 | |
| //! ``` | |
| //! | |
| //! The **home slot** is where the word would like to live. The **tag** is a | |
| //! one-byte summary of the word. Tags are kept in their own row, one byte per | |
| //! slot, beside the row of entries: | |
| //! | |
| //! ```text | |
| //! slot index: 0 1 2 3 4 5 6 7 ... | |
| //! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ | |
| //! tags │ · │ 9B │ · │ · │ C4 │ A7 │ 31 │ A7 │ 1 byte each | |
| //! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ | |
| //! entries │ │"of" │ │ │"cat"│"the"│"sat"│"hat"│ 32 bytes each | |
| //! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ | |
| //! ▲ | |
| //! home slot of "the", | |
| //! and of "hat" as well | |
| //! | |
| //! · = empty slot | |
| //! ``` | |
| //! | |
| //! A word does not always get its home slot: another word may have taken it | |
| //! first, which is what happened to "hat" above. So the word may be in its home | |
| //! slot or in any of the fifteen slots after it. Those sixteen slots are the | |
| //! word's **window**, and the word is either in there or nowhere | |
| //! ([`WALK_WINDOW`]). | |
| //! | |
| //! # Looking a word up | |
| //! | |
| //! [`WordCache::lookup`] walks the window one tag at a time, starting at the home | |
| //! slot, and answers both of a lookup's questions on the way: which slot holds | |
| //! this word, and where would this word go if none of them does. | |
| //! | |
| //! ```text | |
| //! looking up "hat" (tag A7, home slot 5) | |
| //! | |
| //! slot: 5 6 7 8 9 10 11 12 ... | |
| //! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ | |
| //! tags │ A7 │ 31 │ A7 │ · │ 5F │ A7 │ · │ · │ | |
| //! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ | |
| //! entries │"the"│"sat"│"hat"│ │"on" │"mat"│ │ │ | |
| //! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ | |
| //! ▲ ▲ ▲ ▲ | |
| //! │ │ │ └ tagged A7 too, but past the | |
| //! │ │ │ empty slot: out of reach | |
| //! │ │ └ empty: the search stops here | |
| //! │ └ tag matches, and the key says "hat": found it | |
| //! └ tag matches, but the key says "the", so keep looking | |
| //! ``` | |
| //! | |
| //! The walk reads tags, not entries: ruling a slot out costs one byte of memory | |
| //! traffic rather than the 32 an entry takes. It rarely gets far. On running text a | |
| //! word is usually in its home slot, or the slot after it, or not in the table at | |
| //! all. | |
| //! | |
| //! A matching tag is a hint, not an answer. A tag is only one byte of the hash, so | |
| //! about one slot in 255 matches a word that is not there, which is exactly what | |
| //! slot 5 does above. Each matching slot is therefore read and its **key** | |
| //! compared. The key is the copy of the word kept inside the slot, and it is what | |
| //! settles the question. | |
| //! | |
| //! The search stops at the first empty slot. That is safe because storing stops | |
| //! there too: a word is always put in the first empty slot of its window, so no | |
| //! word can ever sit past one. This is why slot 10 is out of reach: whatever is | |
| //! in it, it is not "hat", because "hat" would have taken slot 8. | |
| //! | |
| //! A lookup ends one of two ways ([`Lookup`]): | |
| //! | |
| //! - **Hit**: the ids, ready to use. Nothing is written: a hit leaves the table | |
| //! exactly as it was. | |
| //! - **Miss**: a [`Placement`], the slot this word should take. The caller runs | |
| //! the model, then hands the placement back to [`WordCache::insert`] along with | |
| //! the ids. The placement carries the hash work the lookup already did, so | |
| //! storing the word costs no second hash and no second read of the tags. | |
| //! | |
| //! # Storing a word | |
| //! | |
| //! [`WordCache::insert`] takes the placement from the miss and the ids the model | |
| //! produced, and does four things: | |
| //! | |
| //! 1. **Builds the entry.** The word and its ids go in the slot if they fit; what | |
| //! does not fit is copied into an overflow buffer (below). If those buffers are | |
| //! full, the insert is dropped and nothing changes, since the cache is allowed | |
| //! to forget. | |
| //! 2. **Clears the old entry, if there was one.** When the window had no empty | |
| //! slot the placement points at the home slot, and that word's ids are lost. | |
| //! Any overflow space it was using is handed back. | |
| //! 3. **Writes the tag**, which is what makes the slot findable. | |
| //! 4. **Writes the entry.** | |
| //! | |
| //! # What a slot holds | |
| //! | |
| //! A slot is 32 bytes: sixteen for the key, twelve for the ids, one that says how | |
| //! many ids there are, and three of padding. It comes in three shapes | |
| //! ([`CachedWord`]). The first is the common case, and the reason the sizes are what | |
| //! they are: the whole answer (word and ids) is right there, so a hit reads one | |
| //! slot and stops. | |
| //! | |
| //! ```text | |
| //! ├───── key ─────┤├─── ids ───┤├─────┤ | |
| //! how many | |
| //! | |
| //! "the" → [464] ┌────────────────┬────────────┬─────┐ | |
| //! short word, up to 3 ids │ t h e 3 │ 464 │ 1 │ | |
| //! └────────────────┴────────────┴─────┘ | |
| //! the word, then the ids | |
| //! its length themselves | |
| //! | |
| //! "hello" → [40,71,12,9,33] ┌────────────────┬────────────┬─────┐ | |
| //! short word, too many ids │ h e l l o 5 │ ids are at │ ▒ │ | |
| //! └────────────────┴────────────┴─────┘ | |
| //! | |
| //! "counterrevolutionary" → … ┌────────────────┬────────────┬─────┐ | |
| //! word too long for the key │ hash of word H │ both are at│ ▒ │ | |
| //! └────────────────┴────────────┴─────┘ | |
| //! H says "a hash, where to look in | |
| //! not the word" the buffers | |
| //! | |
| //! ▒ = "not in the slot, follow the offsets" | |
| //! ``` | |
| //! | |
| //! A word of up to 15 bytes is kept whole, with its length in the sixteenth byte. | |
| //! A longer word cannot fit, so the slot keeps the word's hash instead and the | |
| //! bytes go to a buffer. Two different long words can hash to the same number, so | |
| //! for those, and only those, a match is confirmed by comparing the bytes. | |
| //! | |
| //! The two overflow buffers ([`Arena`]) grow up to a budget and then stop. When an | |
| //! entry is evicted, the space it used goes on a free list for its exact | |
| //! length, so the next word of that shape reuses it. Nothing is ever compacted or | |
| //! moved. Words longer than [`MAX_WORD_BYTES`] are not cached at all. | |
| //! | |
| //! # Which entry gets evicted | |
| //! | |
| //! When a word's window is full, one of its entries gets evicted, and it is whatever | |
| //! sits in the home slot. Nothing is measured and nothing is ranked: the newcomer | |
| //! takes the slot it wanted in the first place. | |
| //! | |
| //! A cleverer rule is possible, and this module used to have one: a use counter per | |
| //! entry, the least-used slot of the window loses it, and the counters fade so that | |
| //! a word that was busy early in a document cannot keep its slot forever. Measured | |
| //! against the plain rule, it did keep slightly more of the words worth keeping, and | |
| //! it lost more than it gained to the counter that every hit had to write, so it | |
| //! went. Two things let the plain rule get away with it: nothing is evicted until | |
| //! all sixteen slots of a window are taken, and words repeat closely enough that a | |
| //! word's next use usually comes round before anything has had the chance to evict | |
| //! it. | |
| //! | |
| //! # Why it is built this way | |
| //! | |
| //! - **Tags live in their own row** so that a walk over sixteen slots reads sixteen | |
| //! bytes, which is one or two cache lines. The entries are 32 bytes each, so the | |
| //! same sixteen slots would be 512 bytes and eight cache lines. | |
| //! - **A tag is one byte of the hash**, so a slot that cannot hold the word is ruled | |
| //! out without reading it. All eight bits carry hash: emptiness needs a value no | |
| //! live tag can take rather than a bit of its own, and one value is cheaper than | |
| //! one bit ([`make_tag`]). | |
| //! - **The tag uses the top bits of the hash** because the bottom bits already | |
| //! chose the home slot, and every slot in a window would share those. | |
| //! - **A window is sixteen slots**: how far a word may end up from its home slot | |
| //! before the cache stops looking and evicts instead. Long enough that a full | |
| //! window is uncommon, short enough that walking a full one stays cheap. | |
| //! - **The key is 128 bits** because 15 bytes of word plus one byte of length | |
| //! covers nearly every word whole, and comparing it is then one comparison of one | |
| //! value instead of following a pointer to bytes elsewhere in memory. | |
| //! - **The length is part of the key**, or `"a"` and `"a\0"` would look identical. | |
| //! - **Up to three ids fit in the slot** because most words encode to one, two or | |
| //! three, and that makes a hit a single read. | |
| //! - **The number of slots is a power of two**, so picking the home slot is one | |
| //! bit operation rather than a division. So is folding a step of the walk back | |
| //! into the table when a window runs off the end. | |
| //! - **A miss hands back a placement** because the lookup already knows the hash, | |
| //! the tag and the slot to use, and making the insert work them out again would | |
| //! double that cost on every new word. | |
| //! - **Freed overflow space is filed by exact length**: every run is at most | |
| //! [`MAX_WORD_BYTES`] long, so there can be one free list per length and a freed | |
| //! run always fits the next word of that shape exactly. | |
| //! | |
| //! # Where the ideas come from | |
| //! | |
| //! - [Swiss Tables] is where the control byte comes from: one byte per slot, holding | |
| //! a flag bit and seven bits of hash. Here the whole byte is hash, because the flag | |
| //! bit was there to be tested sixteen lanes at a time and this walk tests one. | |
| //! - [gigatoken] is a BPE tokenizer with a pre-token cache built from the same | |
| //! parts: `u128` packed keys with the length in the top byte, self-contained | |
| //! 32-byte entries, ids inline. It never evicts (it doubles at 3/4 load) and | |
| //! leans on huge pages and prefetching, because its table is sized for DRAM | |
| //! rather than for a CPU cache. | |
| //! - [TinyLFU] (Einziger, Friedman & Manes) is the use-counter-and-fade rule this | |
| //! module tried for eviction and dropped, as its reference implementation | |
| //! [Caffeine] does it. Kept here for whoever wants to try it again. | |
| //! - [huggingface/tokenizers#2234] is an open-addressed cache for this same encode | |
| //! pipeline, arrived at in parallel, fused into the pre-tokenizer's split loop. | |
| //! | |
| //! [Swiss Tables]: https://abseil.io/about/design/swisstables | |
| //! [gigatoken]: https://github.com/marcelroed/gigatoken | |
| //! [TinyLFU]: https://arxiv.org/abs/1512.00727 | |
| //! [Caffeine]: https://github.com/ben-manes/caffeine/blob/master/caffeine/src/main/java/com/github/benmanes/caffeine/cache/FrequencySketch.java | |
| //! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 | |
| //! A table that remembers which token ids a word encodes to, so a model only has to | |
| //! work it out once. | |
| //! | |
| //! Pre-tokenization cuts text into pieces this module calls *words*. Running a model | |
| //! (BPE, Unigram, WordPiece) over one is expensive; words repeat, and a word always | |
| //! encodes to the same ids. So encode each word once and look the answer up after that. | |
| //! | |
| //! # The table | |
| //! | |
| //! Open addressing, linear probing, a fixed power-of-two number of slots, one word per | |
| //! slot. One hash gives both halves of an address, and both rows use it: | |
| //! | |
| //! ```text | |
| //! hash(key) | |
| //! ┌───────────────┴───────────────┐ | |
| //! top 8 bits bottom log2(slots) bits | |
| //! │ │ | |
| //! tag: A7 index: 5 | |
| //! │ │ | |
| //! └───────────────┬───────────────┘ | |
| //! ▼ | |
| //! index: 0 1 2 3 4 5 6 7 ... | |
| //! ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┬─────┐ | |
| //! tags │ · │ 9B │ · │ · │ C4 │ A7 │ 31 │ A7 │ 1 B each | |
| //! ├─────┼─────┼─────┼─────┼─────┼─────┼─────┼─────┤ | |
| //! slots │ │"of" │ │ │"cat"│"the"│"sat"│"hat"│ 32 B each | |
| //! └─────┴─────┴─────┴─────┴─────┴─────┴─────┴─────┘ | |
| //! ▲ | |
| //! home slot of "the" — and of "hat", | |
| //! which took the next free slot after it | |
| //! | |
| //! · = empty | |
| //! ``` | |
| //! | |
| //! A lookup reads `tags` from the home slot forward, [`WALK_WINDOW`] slots at most, and | |
| //! stops at the first empty one. An insert takes that same first empty slot, so nothing | |
| //! is ever stored past one and a walk cannot step over a word that is present. Nothing | |
| //! is ever deleted, so there are no tombstones. A window with no empty slot evicts its | |
| //! home slot. | |
| //! | |
| //! A tag match is a hint, not an answer: one byte of hash, so about one slot in 255 | |
| //! answers for a word that is not in it. The slot's key settles it. | |
| //! | |
| //! # Decisions that aren't obvious | |
| //! | |
| //! - **Two rows rather than one.** Walking sixteen slots reads sixteen bytes of `tags`, | |
| //! one or two cache lines. The same sixteen entries would be 512 bytes and eight lines. | |
| //! - **Index from the low bits, tag from the top eight.** They cannot overlap: the index | |
| //! needs `log2(slots)` bits, so both fit in a `u64` for any table below 2^56 slots — | |
| //! nothing is given up by taking them from opposite ends. The tag *has* to come from | |
| //! bits the index did not use, or every slot in a window would carry the same tag and | |
| //! the walk would learn nothing from reading it. | |
| //! - **The key is hashed, never used as an address directly.** Packed keys for `"tok1"` | |
| //! and `"tok2"` differ in one byte, and the index takes the low bits while the tag takes | |
| //! the top eight, so a one-byte difference has to reach both ends of a `u64`. | |
| //! - **`EMPTY` is a tag value, not a flag bit.** All eight bits carry hash: giving up one | |
| //! value in 256 is cheaper than one bit in eight. A live tag must never be `EMPTY`, or a | |
| //! walk stops short and every entry after it in the window is lost. | |
| //! - **A word of fifteen bytes or fewer is its own key**, its length in the sixteenth byte | |
| //! of a `u128`. That holds nearly every word whole, comparing it is one register-wide | |
| //! equality rather than a pointer chase, and carrying the length is what keeps `"a"` and | |
| //! `"a\0"` apart. | |
| //! - **A longer word keys on 127 bits of its hash.** Widening that key to the field it | |
| //! sits in, rather than to the 64 bits that placed it, is what makes the word's bytes | |
| //! unnecessary: the key comparison is itself the proof. So there is no word arena, no | |
| //! byte compare on the hot path, a hit is the same work for either shape, and no limit | |
| //! on how long a word may be. Bit 127 marks a hashed key, and a packed one always | |
| //! carries a length of 1..15 in its top byte, so the two spaces cannot meet. | |
| //! - **A slot is 32 bytes, so a hit is one read** — and three inline ids is what is left | |
| //! of it after a 16-byte key and a count, not a tuned number. It covers most words in | |
| //! an alphabetic script and fewer in Chinese or Korean, where a word becomes more ids. | |
| //! - **A miss hands back a [`Placement`]** holding the tag and slot the walk already | |
| //! worked out, so storing the word repeats none of it. | |
| //! - **When the ids buffer fills, the whole table is emptied**, rather than each entry's | |
| //! space being tracked and recycled. Forgetting is free; a free list per length was not. | |
| //! - **The cache may forget; it may not lie.** Every trade-off above spends the first | |
| //! freedom and not the second. The single exception is quantified rather than waved at: | |
| //! two words whose 127-bit keys agree would trade ids, at a rate below that of the | |
| //! machine corrupting them in DRAM on the way past. | |
| //! | |
| //! # Prior art | |
| //! | |
| //! - [SwissTables]: the tag row is its control byte and the index/tag split is its H1/H2. | |
| //! Abseil's byte spends one bit on a flag so it can test sixteen lanes at once; this | |
| //! walk tests one lane, so all eight bits are hash. | |
| //! - [gigatoken]: a BPE pre-token cache from the same parts — `u128` keys with the length | |
| //! in the top byte, self-contained 32-byte entries, ids inline. It never evicts | |
| //! (doubling at 3/4 load) and leans on huge pages and prefetching, its table being sized | |
| //! for DRAM rather than for a CPU cache. | |
| //! - [huggingface/tokenizers#2234]: an open-addressed cache for this same pipeline, | |
| //! arrived at in parallel, fused into the pre-tokenizer's split loop. | |
| //! | |
| //! [SwissTables]: https://abseil.io/about/design/swisstables | |
| //! [gigatoken]: https://github.com/marcelroed/gigatoken | |
| //! [huggingface/tokenizers#2234]: https://github.com/huggingface/tokenizers/pull/2234 |
Comment on lines
+303
to
+353
| use ahash::RandomState; | ||
|
|
||
| // ---------------------------------------------------------------- the cache | ||
|
|
||
| /// Longest word the cache will store, in bytes. | ||
| const MAX_WORD_BYTES: usize = 1024; | ||
|
|
||
| /// Word bytes to token ids. See the module docs for the design. | ||
| pub struct WordCache { | ||
| /// The slots, one word each. Their number is a power of two, so a word's home | ||
| /// slot is just the bottom bits of its hash: | ||
| /// ```text | ||
| /// hash(word) & index_mask | ||
| /// ``` | ||
| cached_words: Box<[CachedWord]>, | ||
|
|
||
| /// Hashes a slot's key: the packed word when it is short enough, the word's | ||
| /// bytes when it is not. See [`WordCache::make_word_key`]. | ||
| hasher: RandomState, | ||
|
|
||
| /// `cached_words.len() - 1`. Masks a hash down to its bottom bits, which give the | ||
| /// word's home slot. | ||
| index_mask: usize, | ||
|
|
||
| /// One tag per slot, at the slot's own index: the top byte of the word's hash | ||
| /// ([`make_tag`]), or [`EMPTY`] when the slot holds nothing. | ||
| /// | ||
| /// A walk reads this row instead of the entries, so ruling a slot out costs one | ||
| /// byte rather than the 32 an entry takes. | ||
| tags: Box<[u8]>, | ||
|
|
||
| /// Holds the word's bytes when they don't fit in a [`CachedWord`] key. | ||
| word_bytes_arena: Arena<u8>, | ||
|
|
||
| /// Holds the word's ids when they don't fit in a [`CachedWord`]. | ||
| token_ids_arena: Arena<u32>, | ||
| } | ||
|
|
||
| impl WordCache { | ||
| /// `capacity` is rounded up to a power of two, and to at least one full window. | ||
| pub fn new(capacity: usize) -> Self { | ||
| let n_slots = capacity.next_power_of_two().max(WALK_WINDOW); | ||
| Self { | ||
| hasher: RandomState::new(), | ||
| cached_words: vec![CachedWord::default(); n_slots].into_boxed_slice(), | ||
| tags: vec![EMPTY; n_slots].into_boxed_slice(), | ||
| word_bytes_arena: Arena::new(n_slots * 48), | ||
| token_ids_arena: Arena::new(n_slots * 16), | ||
| index_mask: n_slots - 1, | ||
| } | ||
| } |
Collaborator
There was a problem hiding this comment.
Fixed-seed hashers instead of a per-instance RandomState, so the field goes and the table is reproducible in benches and tests. Nothing here is a security boundary. MAX_WORD_BYTES goes with the word arena; the two arenas collapse to one Vec<u32> and a budget.
Suggested change
| use ahash::RandomState; | |
| // ---------------------------------------------------------------- the cache | |
| /// Longest word the cache will store, in bytes. | |
| const MAX_WORD_BYTES: usize = 1024; | |
| /// Word bytes to token ids. See the module docs for the design. | |
| pub struct WordCache { | |
| /// The slots, one word each. Their number is a power of two, so a word's home | |
| /// slot is just the bottom bits of its hash: | |
| /// ```text | |
| /// hash(word) & index_mask | |
| /// ``` | |
| cached_words: Box<[CachedWord]>, | |
| /// Hashes a slot's key: the packed word when it is short enough, the word's | |
| /// bytes when it is not. See [`WordCache::make_word_key`]. | |
| hasher: RandomState, | |
| /// `cached_words.len() - 1`. Masks a hash down to its bottom bits, which give the | |
| /// word's home slot. | |
| index_mask: usize, | |
| /// One tag per slot, at the slot's own index: the top byte of the word's hash | |
| /// ([`make_tag`]), or [`EMPTY`] when the slot holds nothing. | |
| /// | |
| /// A walk reads this row instead of the entries, so ruling a slot out costs one | |
| /// byte rather than the 32 an entry takes. | |
| tags: Box<[u8]>, | |
| /// Holds the word's bytes when they don't fit in a [`CachedWord`] key. | |
| word_bytes_arena: Arena<u8>, | |
| /// Holds the word's ids when they don't fit in a [`CachedWord`]. | |
| token_ids_arena: Arena<u32>, | |
| } | |
| impl WordCache { | |
| /// `capacity` is rounded up to a power of two, and to at least one full window. | |
| pub fn new(capacity: usize) -> Self { | |
| let n_slots = capacity.next_power_of_two().max(WALK_WINDOW); | |
| Self { | |
| hasher: RandomState::new(), | |
| cached_words: vec![CachedWord::default(); n_slots].into_boxed_slice(), | |
| tags: vec![EMPTY; n_slots].into_boxed_slice(), | |
| word_bytes_arena: Arena::new(n_slots * 48), | |
| token_ids_arena: Arena::new(n_slots * 16), | |
| index_mask: n_slots - 1, | |
| } | |
| } | |
| use ahash::RandomState; | |
| // ---------------------------------------------------------------- the cache | |
| /// Hashes a key to the 64 bits an address is made of, and a long word's bytes to the half | |
| /// of its key that also places it. | |
| /// | |
| /// Fixed seeds rather than [`RandomState::new`]: nothing here is a security boundary, since | |
| /// the worst a collision-flooder wins is a lower hit rate, and a fixed table makes | |
| /// benchmarks and tests reproducible. Digits of pi, so nothing is up anyone's sleeve. | |
| static PLACEMENT: RandomState = RandomState::with_seeds( | |
| 0x243f_6a88_85a3_08d3, | |
| 0x1319_8a2e_0370_7344, | |
| 0xa409_3822_299f_31d0, | |
| 0x082e_fa98_ec4e_6c89, | |
| ); | |
| /// Hashes a long word's bytes a second time, to fill the half of the key that [`PLACEMENT`] | |
| /// does not reach. Two independent hashes are what make the key stand in for the word. | |
| static CONFIRM: RandomState = RandomState::with_seeds( | |
| 0x4528_21e6_38d0_1377, | |
| 0xbe54_66cf_34e9_0c6c, | |
| 0xc0ac_29b7_c97c_50dd, | |
| 0x3f84_d5b5_b547_0917, | |
| ); | |
| /// Word bytes to token ids. See the module docs for the design. | |
| pub struct WordCache { | |
| /// The slots, one word each. Their number is a power of two, so a word's home | |
| /// slot is just the bottom bits of its hash: | |
| /// ```text | |
| /// hash(word) & index_mask | |
| /// ``` | |
| cached_words: Box<[CachedWord]>, | |
| /// `cached_words.len() - 1`. Masks a hash down to its bottom bits, which give the | |
| /// word's home slot. | |
| index_mask: usize, | |
| /// One tag per slot, at the slot's own index: the top byte of the word's hash | |
| /// ([`make_tag`]), or [`EMPTY`] when the slot holds nothing. | |
| /// | |
| /// A walk reads this row instead of the entries, so ruling a slot out costs one | |
| /// byte rather than the 32 an entry takes. | |
| tags: Box<[u8]>, | |
| /// Ids that did not fit in their slot, appended and never moved. A slot points into | |
| /// this with `[offset, how many]`; see [`CachedWord`]. | |
| spilled_ids: Vec<u32>, | |
| /// Ceiling on `spilled_ids`, not a reservation. Reaching it empties the whole table | |
| /// ([`WordCache::insert`]). | |
| spilled_ids_budget: usize, | |
| } | |
| impl WordCache { | |
| /// `capacity` is rounded up to a power of two, and to at least one full window. | |
| pub fn new(capacity: usize) -> Self { | |
| let n_slots = capacity.next_power_of_two().max(WALK_WINDOW); | |
| Self { | |
| cached_words: vec![CachedWord::default(); n_slots].into_boxed_slice(), | |
| tags: vec![EMPTY; n_slots].into_boxed_slice(), | |
| spilled_ids: Vec::new(), | |
| // Capped so that an offset into it always fits the `u32` a slot keeps it in. | |
| spilled_ids_budget: (n_slots * 16).min(u32::MAX as usize), | |
| index_mask: n_slots - 1, | |
| } | |
| } |
Comment on lines
+355
to
+374
| /// The ids `word` encoded to last time, or the [`Placement`] it should be stored | ||
| /// in once the model has worked them out. See [`Lookup`]. | ||
| pub fn lookup<'c, 'w>(&'c self, word: &'w [u8]) -> Lookup<'c, 'w> { | ||
| if word.len() > MAX_WORD_BYTES { | ||
| return Lookup::Miss(None); | ||
| } | ||
|
|
||
| let (key, hash) = self.make_word_key(word); | ||
| match self.find_word_in_cache(key, hash, word) { | ||
| Walk::Found(index) => { | ||
| let slot = self.cached_words[index]; | ||
| Lookup::Hit(if slot.ids_stored_in_arena() { | ||
| self.token_ids_arena.get(slot.ids_off(), slot.ids_len()) | ||
| } else { | ||
| &self.cached_words[index].word_ids[..slot.inline_id_count as usize] | ||
| }) | ||
| } | ||
| Walk::Absent(placement) => Lookup::Miss(Some(placement)), | ||
| } | ||
| } |
Collaborator
There was a problem hiding this comment.
One shape of hit for both kinds of key.
Suggested change
| /// The ids `word` encoded to last time, or the [`Placement`] it should be stored | |
| /// in once the model has worked them out. See [`Lookup`]. | |
| pub fn lookup<'c, 'w>(&'c self, word: &'w [u8]) -> Lookup<'c, 'w> { | |
| if word.len() > MAX_WORD_BYTES { | |
| return Lookup::Miss(None); | |
| } | |
| let (key, hash) = self.make_word_key(word); | |
| match self.find_word_in_cache(key, hash, word) { | |
| Walk::Found(index) => { | |
| let slot = self.cached_words[index]; | |
| Lookup::Hit(if slot.ids_stored_in_arena() { | |
| self.token_ids_arena.get(slot.ids_off(), slot.ids_len()) | |
| } else { | |
| &self.cached_words[index].word_ids[..slot.inline_id_count as usize] | |
| }) | |
| } | |
| Walk::Absent(placement) => Lookup::Miss(Some(placement)), | |
| } | |
| } | |
| /// The ids `word` encoded to last time, or the [`Placement`] it should be stored in | |
| /// once the model has worked them out. See [`Lookup`]. | |
| pub fn lookup(&self, word: &[u8]) -> Lookup<'_> { | |
| let (key, hash) = make_word_key(word); | |
| match self.find_word_in_cache(key, hash) { | |
| Walk::Found(index) => { | |
| let slot = self.cached_words[index]; | |
| Lookup::Hit(if slot.inline_id_count == SPILLED { | |
| let (off, len) = (slot.word_ids[0] as usize, slot.word_ids[1] as usize); | |
| &self.spilled_ids[off..off + len] | |
| } else { | |
| &self.cached_words[index].word_ids[..slot.inline_id_count as usize] | |
| }) | |
| } | |
| Walk::Absent(placement) => Lookup::Miss(placement), | |
| } | |
| } |
Comment on lines
+376
to
+390
| /// Store `ids` as the encoding of the word that `at` was built for. | ||
| /// | ||
| /// Overwrites whatever the slot held, which only happens when the word's window | ||
| /// was full. The module docs say why the choice does not have to be cleverer than | ||
| /// that. If the arenas have no room for the entry, nothing is stored. | ||
| pub fn insert(&mut self, at: Placement<'_>, ids: impl ExactSizeIterator<Item = u32>) { | ||
| let Some(cached_word) = self.build_entry(at.key, at.word, ids) else { | ||
| return; | ||
| }; | ||
| if self.tags[at.index] != EMPTY { | ||
| self.reclaim(at.index); | ||
| } | ||
| self.tags[at.index] = at.tag; | ||
| self.cached_words[at.index] = cached_word; | ||
| } |
Collaborator
There was a problem hiding this comment.
Flush instead of a free list: when the ids buffer is spent, clear the tag row and start over. Forgetting is free; per-length free lists were not. This insert is dropped along with the rest, because at was chosen against contents that no longer exist -- placing it anyway could strand it behind an empty slot.
Suggested change
| /// Store `ids` as the encoding of the word that `at` was built for. | |
| /// | |
| /// Overwrites whatever the slot held, which only happens when the word's window | |
| /// was full. The module docs say why the choice does not have to be cleverer than | |
| /// that. If the arenas have no room for the entry, nothing is stored. | |
| pub fn insert(&mut self, at: Placement<'_>, ids: impl ExactSizeIterator<Item = u32>) { | |
| let Some(cached_word) = self.build_entry(at.key, at.word, ids) else { | |
| return; | |
| }; | |
| if self.tags[at.index] != EMPTY { | |
| self.reclaim(at.index); | |
| } | |
| self.tags[at.index] = at.tag; | |
| self.cached_words[at.index] = cached_word; | |
| } | |
| /// Store `ids` as the encoding of the word that `at` was built for. | |
| /// | |
| /// Overwrites whatever the slot held, which only happens when the word's window | |
| /// was full. The module docs say why the choice does not have to be cleverer than | |
| /// that. | |
| pub fn insert(&mut self, at: Placement, ids: impl ExactSizeIterator<Item = u32>) { | |
| let ids_len = ids.len(); | |
| if ids_len > MAX_INLINE_IDS && self.spilled_ids.len() + ids_len > self.spilled_ids_budget | |
| { | |
| // The buffer is spent. Empty the table rather than track which slot owns which | |
| // run: a tag is what makes a slot findable, so clearing the row strands every | |
| // entry at once. `at` was chosen against the old contents, so drop this insert | |
| // too and let the word be encoded again next time it turns up. | |
| self.tags.fill(EMPTY); | |
| self.spilled_ids.clear(); | |
| return; | |
| } | |
| let entry = self.build_entry(at.key, ids); | |
| self.tags[at.index] = at.tag; | |
| self.cached_words[at.index] = entry; | |
| } |
Comment on lines
+391
to
+408
|
|
||
| /// The value a slot stores in its key, and the hash that places it. A short | ||
| /// word is its own key; a longer one keys on its hash and has to be confirmed | ||
| /// against [`WordCache::word_bytes_arena`], since two long words can hash alike. | ||
| /// | ||
| /// A packed key is hashed as one `u128`, not as the word's bytes it was built | ||
| /// from. Hashing a slice makes aHash mix the length in and then branch on it | ||
| /// to choose a read width; a `u128` is one fixed-width fold with nothing to | ||
| /// decide. The key already carries the length, so nothing is lost. | ||
| fn make_word_key(&self, word: &[u8]) -> (u128, u64) { | ||
| match pack_word(word) { | ||
| Some(packed) => (packed, self.hasher.hash_one(packed)), | ||
| None => { | ||
| let hash = self.hasher.hash_one(word); | ||
| ((hash as u128) | KEY_IS_HASH, hash) | ||
| } | ||
| } | ||
| } |
Collaborator
There was a problem hiding this comment.
make_word_key no longer touches self, so it moves out of the impl -- see the suggestion further down.
Suggested change
| /// The value a slot stores in its key, and the hash that places it. A short | |
| /// word is its own key; a longer one keys on its hash and has to be confirmed | |
| /// against [`WordCache::word_bytes_arena`], since two long words can hash alike. | |
| /// | |
| /// A packed key is hashed as one `u128`, not as the word's bytes it was built | |
| /// from. Hashing a slice makes aHash mix the length in and then branch on it | |
| /// to choose a read width; a `u128` is one fixed-width fold with nothing to | |
| /// decide. The key already carries the length, so nothing is lost. | |
| fn make_word_key(&self, word: &[u8]) -> (u128, u64) { | |
| match pack_word(word) { | |
| Some(packed) => (packed, self.hasher.hash_one(packed)), | |
| None => { | |
| let hash = self.hasher.hash_one(word); | |
| ((hash as u128) | KEY_IS_HASH, hash) | |
| } | |
| } | |
| } |
Comment on lines
+709
to
+740
| /// How many slots [`WordCache::find_word_in_cache`] walks before it stops looking | ||
| /// and evicts: a word is in the sixteen slots from its home slot on, or nowhere. | ||
| /// | ||
| /// Sixteen one-byte tags cross at most two cache lines, and a walk rarely reads | ||
| /// that many, since it stops at the first empty slot, usually a step or two along. | ||
| const WALK_WINDOW: usize = 16; | ||
|
|
||
| /// The tag of a slot that holds nothing. No live entry ever carries it | ||
| /// ([`make_tag`]), which is what lets a walk stop at the first one it reads, | ||
| /// and lets [`WordCache::insert`] tell a slot it has to clear from one it can write | ||
| /// straight into. | ||
| const EMPTY: u8 = 0; | ||
|
|
||
| /// How many bits of the hash a tag carries. All eight of them: a walk therefore | ||
| /// reads a slot it did not want about once in 255. | ||
| const TAG_BITS: u32 = 8; | ||
|
|
||
| /// The tag a hash gives its slot, which is never [`EMPTY`]. | ||
| /// | ||
| /// The top of the hash, because the bottom of it already chose the home slot: a | ||
| /// tag built from those bits would be the same for every slot the word can reach | ||
| /// and would tell a walk nothing. | ||
| /// | ||
| /// Top bytes of 0 and 1 both come out as 1, since [`EMPTY`] is 0 and no live tag may | ||
| /// be that. A word with either of those top bytes therefore shares its tag with one | ||
| /// other, and pays a wasted entry read twice as often as the rest. A live tag of | ||
| /// [`EMPTY`] would cost far more: a walk would stop at that slot, losing every entry | ||
| /// stored after it in the window, and the insert that overwrote it would hand its | ||
| /// arena runs back to nobody. | ||
| fn make_tag(hash: u64) -> u8 { | ||
| ((hash >> (u64::BITS - TAG_BITS)) as u8).max(EMPTY + 1) | ||
| } |
Collaborator
There was a problem hiding this comment.
TAG_BITS was a named constant for "one byte".
Suggested change
| /// How many slots [`WordCache::find_word_in_cache`] walks before it stops looking | |
| /// and evicts: a word is in the sixteen slots from its home slot on, or nowhere. | |
| /// | |
| /// Sixteen one-byte tags cross at most two cache lines, and a walk rarely reads | |
| /// that many, since it stops at the first empty slot, usually a step or two along. | |
| const WALK_WINDOW: usize = 16; | |
| /// The tag of a slot that holds nothing. No live entry ever carries it | |
| /// ([`make_tag`]), which is what lets a walk stop at the first one it reads, | |
| /// and lets [`WordCache::insert`] tell a slot it has to clear from one it can write | |
| /// straight into. | |
| const EMPTY: u8 = 0; | |
| /// How many bits of the hash a tag carries. All eight of them: a walk therefore | |
| /// reads a slot it did not want about once in 255. | |
| const TAG_BITS: u32 = 8; | |
| /// The tag a hash gives its slot, which is never [`EMPTY`]. | |
| /// | |
| /// The top of the hash, because the bottom of it already chose the home slot: a | |
| /// tag built from those bits would be the same for every slot the word can reach | |
| /// and would tell a walk nothing. | |
| /// | |
| /// Top bytes of 0 and 1 both come out as 1, since [`EMPTY`] is 0 and no live tag may | |
| /// be that. A word with either of those top bytes therefore shares its tag with one | |
| /// other, and pays a wasted entry read twice as often as the rest. A live tag of | |
| /// [`EMPTY`] would cost far more: a walk would stop at that slot, losing every entry | |
| /// stored after it in the window, and the insert that overwrote it would hand its | |
| /// arena runs back to nobody. | |
| fn make_tag(hash: u64) -> u8 { | |
| ((hash >> (u64::BITS - TAG_BITS)) as u8).max(EMPTY + 1) | |
| } | |
| /// How many slots [`WordCache::find_word_in_cache`] walks before it stops looking | |
| /// and evicts: a word is in the sixteen slots from its home slot on, or nowhere. | |
| /// | |
| /// Sixteen one-byte tags cross at most two cache lines, and a walk rarely reads | |
| /// that many, since it stops at the first empty slot, usually a step or two along. | |
| const WALK_WINDOW: usize = 16; | |
| /// The tag of a slot that holds nothing. No live entry ever carries it | |
| /// ([`make_tag`]), which is what lets a walk stop at the first one it reads. | |
| const EMPTY: u8 = 0; | |
| /// The tag a hash gives its slot, which is never [`EMPTY`]. | |
| /// | |
| /// The top byte, because the bottom of the hash already chose the home slot: a tag built | |
| /// from those bits would be the same for every slot the word can reach and would tell a | |
| /// walk nothing. | |
| /// | |
| /// Top bytes of 0 and 1 both come out as 1, since [`EMPTY`] is 0 and no live tag may be | |
| /// that. A word with either of those top bytes therefore shares its tag with one other, and | |
| /// pays a wasted entry read twice as often as the rest. A live tag of [`EMPTY`] would cost | |
| /// far more: a walk would stop at that slot, losing every entry stored after it in the | |
| /// window. | |
| fn make_tag(hash: u64) -> u8 { | |
| ((hash >> 56) as u8).max(EMPTY + 1) | |
| } |
Comment on lines
+742
to
+1027
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
|
|
||
| /// Look a word up and store what it encoded to, for tests that care about | ||
| /// what the table ends up holding rather than about the two steps. | ||
| fn store(cache: &mut WordCache, word: &[u8], ids: impl ExactSizeIterator<Item = u32>) { | ||
| let at = match cache.lookup(word) { | ||
| Lookup::Miss(at) => at, | ||
| Lookup::Hit(_) => None, | ||
| }; | ||
| if let Some(at) = at { | ||
| cache.insert(at, ids); | ||
| } | ||
| } | ||
|
|
||
| /// The slot a word is in, or `None` if the table does not hold it. | ||
| fn slot_of(cache: &WordCache, word: &[u8]) -> Option<usize> { | ||
| let (key, hash) = cache.make_word_key(word); | ||
| match cache.find_word_in_cache(key, hash, word) { | ||
| Walk::Found(index) => Some(index), | ||
| Walk::Absent(_) => None, | ||
| } | ||
| } | ||
|
|
||
| #[test] | ||
| fn roundtrip() { | ||
| let mut cache = WordCache::new(1 << 8); | ||
| assert_eq!(cache.lookup(b"hello").hit(), None); | ||
| store(&mut cache, b"hello", [1u32, 2, 3].into_iter()); | ||
| store(&mut cache, b"world", [4u32].into_iter()); | ||
| assert_eq!(cache.lookup(b"hello").hit(), Some(&[1u32, 2, 3][..])); | ||
| assert_eq!(cache.lookup(b"world").hit(), Some(&[4u32][..])); | ||
| assert_eq!(cache.lookup(b"hell").hit(), None); | ||
| } | ||
|
|
||
| /// The two properties the key encoding rests on: a packed key is unique per | ||
| /// word, and never looks like a hashed one. | ||
| #[test] | ||
| fn packed_keys_are_unique_and_never_look_hashed() { | ||
| assert_ne!(pack_word(b"a"), pack_word(b"a\0")); | ||
| assert_eq!(pack_word(&[0u8; 15]).unwrap() & KEY_IS_HASH, 0); | ||
| assert_eq!(pack_word(b""), None); | ||
| assert_eq!(pack_word(&[b'x'; 16]), None); | ||
| } | ||
|
|
||
| /// A live entry's tag has to be something [`EMPTY`] is not, or a walk reads an | ||
| /// occupied slot as the end of the chain, stores on top of it, and drops the | ||
| /// arena runs it was using on the floor. Every tag is a byte of hash now, so | ||
| /// the hashes whose top byte is zero are the ones that have to be caught. | ||
| #[test] | ||
| fn a_live_tag_is_never_the_empty_marker() { | ||
| for hash in [0u64, 1, u64::MAX, 1 << 57, u64::MAX >> TAG_BITS] { | ||
| assert_ne!(make_tag(hash), EMPTY, "hash {hash:#x}"); | ||
| } | ||
| // Any other top byte is a tag in its own right, kept as it is. | ||
| assert_eq!(make_tag(0xA7 << (u64::BITS - TAG_BITS)), 0xA7); | ||
| } | ||
|
|
||
| /// A short word's placement comes out of its packed key rather than its | ||
| /// bytes, which leaves the hash doing all of the mixing. Words that differ in | ||
| /// one byte pack to keys that differ in one byte, so an index taken from those | ||
| /// bits as they are (`packed as u64`) drops most of these on one slot. | ||
| #[test] | ||
| fn short_words_spread_across_the_table() { | ||
| let cache = WordCache::new(1 << 12); | ||
| let homes: std::collections::HashSet<usize> = (0..1000) | ||
| .map(|i| { | ||
| let (_, hash) = cache.make_word_key(format!("tok{i}").as_bytes()); | ||
| hash as usize & cache.index_mask | ||
| }) | ||
| .collect(); | ||
| // 1000 words over 4096 slots share homes by chance alone; ~887 distinct is | ||
| // as good as a perfect hash gets, so the floor is well under it. | ||
| assert!(homes.len() > 820, "only {} distinct homes", homes.len()); | ||
| } | ||
|
|
||
| #[test] | ||
| fn oversized_words_are_ignored() { | ||
| let mut cache = WordCache::new(1 << 8); | ||
| let big = vec![7u8; MAX_WORD_BYTES + 1]; | ||
| store(&mut cache, &big, [1u32].into_iter()); | ||
| assert_eq!(cache.lookup(&big).hit(), None); | ||
| } | ||
|
|
||
| /// Both sides of the 15-byte packing boundary and both sides of the inline/ | ||
| /// arena boundary have to survive a round trip, including the long word whose | ||
| /// stored `key` is only a hash and needs the byte comparison to confirm it. | ||
| #[test] | ||
| fn every_slot_shape_round_trips() { | ||
| let mut cache = WordCache::new(1 << 8); | ||
| let long = vec![b'x'; 200]; | ||
| let cases: [(&[u8], Vec<u32>); 6] = [ | ||
| (b"short", vec![1]), | ||
| (b"short-wide", (0..40).collect()), | ||
| (b"fifteen-bytes.", vec![2]), | ||
| (b"sixteen-bytes.aa", vec![3]), | ||
| (&long, vec![9]), | ||
| (&long[..64], (0..64).collect()), | ||
| ]; | ||
| for (word, ids) in &cases { | ||
| store(&mut cache, word, ids.clone().into_iter()); | ||
| } | ||
| for (word, ids) in &cases { | ||
| assert_eq!(cache.lookup(word).hit(), Some(&ids[..]), "{word:?}"); | ||
| } | ||
| } | ||
|
|
||
| /// A spilled entry keeps the word's byte count and its id count in eleven bits | ||
| /// each, side by side in one `u32`. The longest word the cache accepts, when it | ||
| /// encodes to one id per byte, is the largest either of them can get. If they | ||
| /// ever overlapped, a hit would read the wrong lengths and hand back the wrong | ||
| /// ids. | ||
| #[test] | ||
| fn a_word_at_the_length_limit_round_trips() { | ||
| let mut cache = WordCache::new(1 << 8); | ||
| let word = vec![b'q'; MAX_WORD_BYTES]; | ||
| let ids: Vec<u32> = (0..MAX_WORD_BYTES as u32).collect(); | ||
| store(&mut cache, &word, ids.clone().into_iter()); | ||
| assert_eq!(cache.lookup(&word).hit(), Some(&ids[..])); | ||
| } | ||
|
|
||
| /// A tag is one byte, so one slot in 255 answers for a word that is not in | ||
| /// it. That is too rare to wait for, so forge one: put a word's tag on a slot | ||
| /// holding a different key, and demand the walk look past it rather than | ||
| /// treat the tag as the answer. | ||
| #[test] | ||
| fn a_tag_collision_is_confirmed_against_the_key() { | ||
| let mut cache = WordCache::new(WALK_WINDOW); | ||
| let (_, hash) = cache.make_word_key(b"beta"); | ||
| let decoy = hash as usize & cache.index_mask; | ||
| cache.tags[decoy] = make_tag(hash); | ||
| // Any key that is not beta's. A packed key always carries a length in its | ||
| // top byte, so 1 cannot be one. | ||
| cache.cached_words[decoy] = CachedWord { | ||
| word_bytes_or_hash: 1, | ||
| word_ids: [7, 0, 0], | ||
| inline_id_count: 1, | ||
| }; | ||
|
|
||
| store(&mut cache, b"beta", [2u32].into_iter()); | ||
| assert_eq!(cache.lookup(b"beta").hit(), Some(&[2u32][..])); | ||
| } | ||
|
|
||
| /// Two long words can hash to the same key, and then only their bytes tell | ||
| /// them apart. Real hash collisions are too rare to write a test around, so | ||
| /// forge one: park another word's entry on this word's home slot, stamp this | ||
| /// word's key and tag on it, and demand the walk look past it. | ||
| #[test] | ||
| fn a_hashed_key_is_confirmed_against_the_word_bytes() { | ||
| let mut cache = WordCache::new(1 << 8); | ||
| let mine = vec![b'a'; 40]; | ||
| let theirs = vec![b'b'; 40]; | ||
|
|
||
| store(&mut cache, &theirs, [7u32].into_iter()); | ||
| let their_index = slot_of(&cache, &theirs).unwrap(); | ||
| let their_slot = cache.cached_words[their_index]; | ||
| let (my_key, my_hash) = cache.make_word_key(&mine); | ||
| let my_home = my_hash as usize & cache.index_mask; | ||
| cache.cached_words[my_home] = CachedWord { | ||
| word_bytes_or_hash: my_key, | ||
| ..their_slot | ||
| }; | ||
| cache.tags[my_home] = make_tag(my_hash); | ||
|
|
||
| store(&mut cache, &mine, [1u32, 2].into_iter()); | ||
| assert_eq!(cache.lookup(&mine).hit(), Some(&[1u32, 2][..])); | ||
| } | ||
|
|
||
| /// A word that hashes into a full window is stored rather than turned away, | ||
| /// and the entry evicted is the one in its home slot, even when that entry has | ||
| /// just been used, which is the whole of what this policy costs. | ||
| #[test] | ||
| fn a_full_window_evicts_its_home_slot() { | ||
| let mut cache = WordCache::new(WALK_WINDOW); | ||
| let words: Vec<Vec<u8>> = (0..WALK_WINDOW as u8).map(|i| vec![i; 4]).collect(); | ||
| for (i, word) in words.iter().enumerate() { | ||
| store(&mut cache, word, [i as u32].into_iter()); | ||
| } | ||
| let (_, hash) = cache.make_word_key(b"newcomer"); | ||
| let home = hash as usize & cache.index_mask; | ||
| let evicted = words | ||
| .iter() | ||
| .position(|word| slot_of(&cache, word) == Some(home)) | ||
| .expect("the table is full, so some word holds that slot"); | ||
| for _ in 0..8 { | ||
| assert!(cache.lookup(&words[evicted]).hit().is_some()); | ||
| } | ||
|
|
||
| store(&mut cache, b"newcomer", [999u32].into_iter()); | ||
| assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999u32][..])); | ||
| assert_eq!(cache.lookup(&words[evicted]).hit(), None); | ||
| for (i, word) in words.iter().enumerate() { | ||
| if i != evicted { | ||
| assert_eq!( | ||
| cache.lookup(word).hit(), | ||
| Some(&[i as u32][..]), | ||
| "entry {i} was dropped as well" | ||
| ); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A window that runs off the end of the table carries on at the start, and a | ||
| /// word stored past that seam has to be found there. Nothing mirrors the first | ||
| /// tags at the end of the row. Every step of the walk is folded back into the | ||
| /// table instead. | ||
| #[test] | ||
| fn a_word_stored_past_the_end_of_the_table_is_found() { | ||
| let mut cache = WordCache::new(WALK_WINDOW); | ||
| let last_slot = cache.cached_words.len() - 1; | ||
| let homed_on_the_last_slot: Vec<Vec<u8>> = (0..4000u32) | ||
| .map(|i| format!("w{i}").into_bytes()) | ||
| .filter(|word| { | ||
| let (_, hash) = cache.make_word_key(word); | ||
| hash as usize & cache.index_mask == last_slot | ||
| }) | ||
| .take(3) | ||
| .collect(); | ||
| assert_eq!(homed_on_the_last_slot.len(), 3); | ||
|
|
||
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | ||
| store(&mut cache, word, [i as u32].into_iter()); | ||
| } | ||
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | ||
| assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); | ||
| } | ||
| // The first one took the last slot, so the other two could only go past it. | ||
| assert!(slot_of(&cache, &homed_on_the_last_slot[1]).unwrap() < last_slot); | ||
| } | ||
|
|
||
| /// Reusing an evicted entry's arena run is where this design can go wrong: | ||
| /// hand one run to two live entries and a hit starts returning another word's | ||
| /// ids. Churn a table far too small for the input and demand the invariant | ||
| /// that matters: an entry may be evicted, but a hit is never wrong. | ||
| #[test] | ||
| fn a_hit_never_returns_another_words_ids() { | ||
| let mut cache = WordCache::new(64); | ||
| let mut expected: Vec<(Vec<u8>, Vec<u32>)> = Vec::new(); | ||
| for i in 0..2000usize { | ||
| let word = match i % 4 { | ||
| 0 => format!("w{i}"), | ||
| 1 => format!("a-long-word-past-fifteen-bytes-{i}"), | ||
| 2 => format!("k{i}xxxxxxxxxxxx"), | ||
| _ => format!("{}-{i}", "z".repeat(i % 40)), | ||
| } | ||
| .into_bytes(); | ||
| let ids: Vec<u32> = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); | ||
| store(&mut cache, &word, ids.clone().into_iter()); | ||
| expected.push((word, ids)); | ||
| } | ||
| let mut live = 0; | ||
| for (word, ids) in &expected { | ||
| if let Some(hit) = cache.lookup(word).hit() { | ||
| assert_eq!(hit, &ids[..], "{word:?}"); | ||
| live += 1; | ||
| } | ||
| } | ||
| assert!( | ||
| live > 0, | ||
| "everything was evicted, so the test proves nothing" | ||
| ); | ||
| } | ||
|
|
||
| /// Freed runs have to come back. Without reuse the arenas grow with every | ||
| /// insert until their budget is spent, and from then on the cache silently | ||
| /// stops accepting words even though the table has room for them. | ||
| #[test] | ||
| fn arenas_hold_the_live_set_not_every_insert() { | ||
| let mut cache = WordCache::new(64); | ||
| // Same length every time, so one free list serves every word and the | ||
| // bounds below are exact: 64 slots, plus the run reserved for the insert | ||
| // in flight before the entry it replaces gives its own run back. | ||
| let word = |i: usize| format!("a-long-word-past-fifteen-bytes-{i:04}"); | ||
| for i in 0..5000usize { | ||
| store(&mut cache, word(i).as_bytes(), [i as u32; 8].into_iter()); | ||
| } | ||
| assert_eq!( | ||
| cache.lookup(word(4999).as_bytes()).hit(), | ||
| Some(&[4999u32; 8][..]), | ||
| "inserts stopped landing: the arenas ran out" | ||
| ); | ||
| assert!(cache.word_bytes_arena.data.len() <= 65 * 35); | ||
| assert!(cache.token_ids_arena.data.len() <= 65 * 8); | ||
| } | ||
| } |
Collaborator
There was a problem hiding this comment.
Tests, as one block since the edits are spread thin. Gone: oversized_words_are_ignored (no cap left) and a_hashed_key_is_confirmed_against_the_word_bytes (no byte compare left). a_word_at_the_length_limit_round_trips becomes a_very_long_word_round_trips at 8 KiB, and the arena test becomes one that the flush has to pass. a_full_window_evicts_its_home_slot keeps the half that could actually break -- the neighbours surviving.
Suggested change
| #[cfg(test)] | |
| mod tests { | |
| use super::*; | |
| /// Look a word up and store what it encoded to, for tests that care about | |
| /// what the table ends up holding rather than about the two steps. | |
| fn store(cache: &mut WordCache, word: &[u8], ids: impl ExactSizeIterator<Item = u32>) { | |
| let at = match cache.lookup(word) { | |
| Lookup::Miss(at) => at, | |
| Lookup::Hit(_) => None, | |
| }; | |
| if let Some(at) = at { | |
| cache.insert(at, ids); | |
| } | |
| } | |
| /// The slot a word is in, or `None` if the table does not hold it. | |
| fn slot_of(cache: &WordCache, word: &[u8]) -> Option<usize> { | |
| let (key, hash) = cache.make_word_key(word); | |
| match cache.find_word_in_cache(key, hash, word) { | |
| Walk::Found(index) => Some(index), | |
| Walk::Absent(_) => None, | |
| } | |
| } | |
| #[test] | |
| fn roundtrip() { | |
| let mut cache = WordCache::new(1 << 8); | |
| assert_eq!(cache.lookup(b"hello").hit(), None); | |
| store(&mut cache, b"hello", [1u32, 2, 3].into_iter()); | |
| store(&mut cache, b"world", [4u32].into_iter()); | |
| assert_eq!(cache.lookup(b"hello").hit(), Some(&[1u32, 2, 3][..])); | |
| assert_eq!(cache.lookup(b"world").hit(), Some(&[4u32][..])); | |
| assert_eq!(cache.lookup(b"hell").hit(), None); | |
| } | |
| /// The two properties the key encoding rests on: a packed key is unique per | |
| /// word, and never looks like a hashed one. | |
| #[test] | |
| fn packed_keys_are_unique_and_never_look_hashed() { | |
| assert_ne!(pack_word(b"a"), pack_word(b"a\0")); | |
| assert_eq!(pack_word(&[0u8; 15]).unwrap() & KEY_IS_HASH, 0); | |
| assert_eq!(pack_word(b""), None); | |
| assert_eq!(pack_word(&[b'x'; 16]), None); | |
| } | |
| /// A live entry's tag has to be something [`EMPTY`] is not, or a walk reads an | |
| /// occupied slot as the end of the chain, stores on top of it, and drops the | |
| /// arena runs it was using on the floor. Every tag is a byte of hash now, so | |
| /// the hashes whose top byte is zero are the ones that have to be caught. | |
| #[test] | |
| fn a_live_tag_is_never_the_empty_marker() { | |
| for hash in [0u64, 1, u64::MAX, 1 << 57, u64::MAX >> TAG_BITS] { | |
| assert_ne!(make_tag(hash), EMPTY, "hash {hash:#x}"); | |
| } | |
| // Any other top byte is a tag in its own right, kept as it is. | |
| assert_eq!(make_tag(0xA7 << (u64::BITS - TAG_BITS)), 0xA7); | |
| } | |
| /// A short word's placement comes out of its packed key rather than its | |
| /// bytes, which leaves the hash doing all of the mixing. Words that differ in | |
| /// one byte pack to keys that differ in one byte, so an index taken from those | |
| /// bits as they are (`packed as u64`) drops most of these on one slot. | |
| #[test] | |
| fn short_words_spread_across_the_table() { | |
| let cache = WordCache::new(1 << 12); | |
| let homes: std::collections::HashSet<usize> = (0..1000) | |
| .map(|i| { | |
| let (_, hash) = cache.make_word_key(format!("tok{i}").as_bytes()); | |
| hash as usize & cache.index_mask | |
| }) | |
| .collect(); | |
| // 1000 words over 4096 slots share homes by chance alone; ~887 distinct is | |
| // as good as a perfect hash gets, so the floor is well under it. | |
| assert!(homes.len() > 820, "only {} distinct homes", homes.len()); | |
| } | |
| #[test] | |
| fn oversized_words_are_ignored() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let big = vec![7u8; MAX_WORD_BYTES + 1]; | |
| store(&mut cache, &big, [1u32].into_iter()); | |
| assert_eq!(cache.lookup(&big).hit(), None); | |
| } | |
| /// Both sides of the 15-byte packing boundary and both sides of the inline/ | |
| /// arena boundary have to survive a round trip, including the long word whose | |
| /// stored `key` is only a hash and needs the byte comparison to confirm it. | |
| #[test] | |
| fn every_slot_shape_round_trips() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let long = vec![b'x'; 200]; | |
| let cases: [(&[u8], Vec<u32>); 6] = [ | |
| (b"short", vec![1]), | |
| (b"short-wide", (0..40).collect()), | |
| (b"fifteen-bytes.", vec![2]), | |
| (b"sixteen-bytes.aa", vec![3]), | |
| (&long, vec![9]), | |
| (&long[..64], (0..64).collect()), | |
| ]; | |
| for (word, ids) in &cases { | |
| store(&mut cache, word, ids.clone().into_iter()); | |
| } | |
| for (word, ids) in &cases { | |
| assert_eq!(cache.lookup(word).hit(), Some(&ids[..]), "{word:?}"); | |
| } | |
| } | |
| /// A spilled entry keeps the word's byte count and its id count in eleven bits | |
| /// each, side by side in one `u32`. The longest word the cache accepts, when it | |
| /// encodes to one id per byte, is the largest either of them can get. If they | |
| /// ever overlapped, a hit would read the wrong lengths and hand back the wrong | |
| /// ids. | |
| #[test] | |
| fn a_word_at_the_length_limit_round_trips() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let word = vec![b'q'; MAX_WORD_BYTES]; | |
| let ids: Vec<u32> = (0..MAX_WORD_BYTES as u32).collect(); | |
| store(&mut cache, &word, ids.clone().into_iter()); | |
| assert_eq!(cache.lookup(&word).hit(), Some(&ids[..])); | |
| } | |
| /// A tag is one byte, so one slot in 255 answers for a word that is not in | |
| /// it. That is too rare to wait for, so forge one: put a word's tag on a slot | |
| /// holding a different key, and demand the walk look past it rather than | |
| /// treat the tag as the answer. | |
| #[test] | |
| fn a_tag_collision_is_confirmed_against_the_key() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let (_, hash) = cache.make_word_key(b"beta"); | |
| let decoy = hash as usize & cache.index_mask; | |
| cache.tags[decoy] = make_tag(hash); | |
| // Any key that is not beta's. A packed key always carries a length in its | |
| // top byte, so 1 cannot be one. | |
| cache.cached_words[decoy] = CachedWord { | |
| word_bytes_or_hash: 1, | |
| word_ids: [7, 0, 0], | |
| inline_id_count: 1, | |
| }; | |
| store(&mut cache, b"beta", [2u32].into_iter()); | |
| assert_eq!(cache.lookup(b"beta").hit(), Some(&[2u32][..])); | |
| } | |
| /// Two long words can hash to the same key, and then only their bytes tell | |
| /// them apart. Real hash collisions are too rare to write a test around, so | |
| /// forge one: park another word's entry on this word's home slot, stamp this | |
| /// word's key and tag on it, and demand the walk look past it. | |
| #[test] | |
| fn a_hashed_key_is_confirmed_against_the_word_bytes() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let mine = vec![b'a'; 40]; | |
| let theirs = vec![b'b'; 40]; | |
| store(&mut cache, &theirs, [7u32].into_iter()); | |
| let their_index = slot_of(&cache, &theirs).unwrap(); | |
| let their_slot = cache.cached_words[their_index]; | |
| let (my_key, my_hash) = cache.make_word_key(&mine); | |
| let my_home = my_hash as usize & cache.index_mask; | |
| cache.cached_words[my_home] = CachedWord { | |
| word_bytes_or_hash: my_key, | |
| ..their_slot | |
| }; | |
| cache.tags[my_home] = make_tag(my_hash); | |
| store(&mut cache, &mine, [1u32, 2].into_iter()); | |
| assert_eq!(cache.lookup(&mine).hit(), Some(&[1u32, 2][..])); | |
| } | |
| /// A word that hashes into a full window is stored rather than turned away, | |
| /// and the entry evicted is the one in its home slot, even when that entry has | |
| /// just been used, which is the whole of what this policy costs. | |
| #[test] | |
| fn a_full_window_evicts_its_home_slot() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let words: Vec<Vec<u8>> = (0..WALK_WINDOW as u8).map(|i| vec![i; 4]).collect(); | |
| for (i, word) in words.iter().enumerate() { | |
| store(&mut cache, word, [i as u32].into_iter()); | |
| } | |
| let (_, hash) = cache.make_word_key(b"newcomer"); | |
| let home = hash as usize & cache.index_mask; | |
| let evicted = words | |
| .iter() | |
| .position(|word| slot_of(&cache, word) == Some(home)) | |
| .expect("the table is full, so some word holds that slot"); | |
| for _ in 0..8 { | |
| assert!(cache.lookup(&words[evicted]).hit().is_some()); | |
| } | |
| store(&mut cache, b"newcomer", [999u32].into_iter()); | |
| assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999u32][..])); | |
| assert_eq!(cache.lookup(&words[evicted]).hit(), None); | |
| for (i, word) in words.iter().enumerate() { | |
| if i != evicted { | |
| assert_eq!( | |
| cache.lookup(word).hit(), | |
| Some(&[i as u32][..]), | |
| "entry {i} was dropped as well" | |
| ); | |
| } | |
| } | |
| } | |
| /// A window that runs off the end of the table carries on at the start, and a | |
| /// word stored past that seam has to be found there. Nothing mirrors the first | |
| /// tags at the end of the row. Every step of the walk is folded back into the | |
| /// table instead. | |
| #[test] | |
| fn a_word_stored_past_the_end_of_the_table_is_found() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let last_slot = cache.cached_words.len() - 1; | |
| let homed_on_the_last_slot: Vec<Vec<u8>> = (0..4000u32) | |
| .map(|i| format!("w{i}").into_bytes()) | |
| .filter(|word| { | |
| let (_, hash) = cache.make_word_key(word); | |
| hash as usize & cache.index_mask == last_slot | |
| }) | |
| .take(3) | |
| .collect(); | |
| assert_eq!(homed_on_the_last_slot.len(), 3); | |
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | |
| store(&mut cache, word, [i as u32].into_iter()); | |
| } | |
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | |
| assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); | |
| } | |
| // The first one took the last slot, so the other two could only go past it. | |
| assert!(slot_of(&cache, &homed_on_the_last_slot[1]).unwrap() < last_slot); | |
| } | |
| /// Reusing an evicted entry's arena run is where this design can go wrong: | |
| /// hand one run to two live entries and a hit starts returning another word's | |
| /// ids. Churn a table far too small for the input and demand the invariant | |
| /// that matters: an entry may be evicted, but a hit is never wrong. | |
| #[test] | |
| fn a_hit_never_returns_another_words_ids() { | |
| let mut cache = WordCache::new(64); | |
| let mut expected: Vec<(Vec<u8>, Vec<u32>)> = Vec::new(); | |
| for i in 0..2000usize { | |
| let word = match i % 4 { | |
| 0 => format!("w{i}"), | |
| 1 => format!("a-long-word-past-fifteen-bytes-{i}"), | |
| 2 => format!("k{i}xxxxxxxxxxxx"), | |
| _ => format!("{}-{i}", "z".repeat(i % 40)), | |
| } | |
| .into_bytes(); | |
| let ids: Vec<u32> = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); | |
| store(&mut cache, &word, ids.clone().into_iter()); | |
| expected.push((word, ids)); | |
| } | |
| let mut live = 0; | |
| for (word, ids) in &expected { | |
| if let Some(hit) = cache.lookup(word).hit() { | |
| assert_eq!(hit, &ids[..], "{word:?}"); | |
| live += 1; | |
| } | |
| } | |
| assert!( | |
| live > 0, | |
| "everything was evicted, so the test proves nothing" | |
| ); | |
| } | |
| /// Freed runs have to come back. Without reuse the arenas grow with every | |
| /// insert until their budget is spent, and from then on the cache silently | |
| /// stops accepting words even though the table has room for them. | |
| #[test] | |
| fn arenas_hold_the_live_set_not_every_insert() { | |
| let mut cache = WordCache::new(64); | |
| // Same length every time, so one free list serves every word and the | |
| // bounds below are exact: 64 slots, plus the run reserved for the insert | |
| // in flight before the entry it replaces gives its own run back. | |
| let word = |i: usize| format!("a-long-word-past-fifteen-bytes-{i:04}"); | |
| for i in 0..5000usize { | |
| store(&mut cache, word(i).as_bytes(), [i as u32; 8].into_iter()); | |
| } | |
| assert_eq!( | |
| cache.lookup(word(4999).as_bytes()).hit(), | |
| Some(&[4999u32; 8][..]), | |
| "inserts stopped landing: the arenas ran out" | |
| ); | |
| assert!(cache.word_bytes_arena.data.len() <= 65 * 35); | |
| assert!(cache.token_ids_arena.data.len() <= 65 * 8); | |
| } | |
| } | |
| #[cfg(test)] | |
| mod tests { | |
| use super::*; | |
| /// Look a word up and store what it encoded to, for tests that care about | |
| /// what the table ends up holding rather than about the two steps. | |
| fn store(cache: &mut WordCache, word: &[u8], ids: impl ExactSizeIterator<Item = u32>) { | |
| if let Lookup::Miss(at) = cache.lookup(word) { | |
| cache.insert(at, ids); | |
| } | |
| } | |
| /// The slot a word is in, or `None` if the table does not hold it. | |
| fn slot_of(cache: &WordCache, word: &[u8]) -> Option<usize> { | |
| let (key, hash) = make_word_key(word); | |
| match cache.find_word_in_cache(key, hash) { | |
| Walk::Found(index) => Some(index), | |
| Walk::Absent(_) => None, | |
| } | |
| } | |
| #[test] | |
| fn roundtrip() { | |
| let mut cache = WordCache::new(1 << 8); | |
| assert_eq!(cache.lookup(b"hello").hit(), None); | |
| store(&mut cache, b"hello", [1u32, 2, 3].into_iter()); | |
| store(&mut cache, b"world", [4u32].into_iter()); | |
| assert_eq!(cache.lookup(b"hello").hit(), Some(&[1u32, 2, 3][..])); | |
| assert_eq!(cache.lookup(b"world").hit(), Some(&[4u32][..])); | |
| assert_eq!(cache.lookup(b"hell").hit(), None); | |
| } | |
| /// The two properties the key encoding rests on: a packed key is unique per | |
| /// word, and never looks like a hashed one. | |
| #[test] | |
| fn packed_keys_are_unique_and_never_look_hashed() { | |
| assert_ne!(pack_word(b"a"), pack_word(b"a\0")); | |
| assert_eq!(pack_word(&[0u8; 15]).unwrap() & KEY_IS_HASH, 0); | |
| assert_eq!(pack_word(b""), None); | |
| assert_eq!(pack_word(&[b'x'; 16]), None); | |
| } | |
| /// A live entry's tag has to be something [`EMPTY`] is not, or a walk reads an | |
| /// occupied slot as the end of the chain and stores on top of it. Every tag is a byte | |
| /// of hash, so the hashes whose top byte is zero are the ones that have to be caught. | |
| #[test] | |
| fn a_live_tag_is_never_the_empty_marker() { | |
| for hash in [0u64, 1, u64::MAX, 1 << 57, u64::MAX >> 8] { | |
| assert_ne!(make_tag(hash), EMPTY, "hash {hash:#x}"); | |
| } | |
| // Any other top byte is a tag in its own right, kept as it is. | |
| assert_eq!(make_tag(0xA7 << 56), 0xA7); | |
| } | |
| /// A short word's placement comes out of its packed key rather than its | |
| /// bytes, which leaves the hash doing all of the mixing. Words that differ in | |
| /// one byte pack to keys that differ in one byte, so an index taken from those | |
| /// bits as they are (`packed as u64`) drops most of these on one slot. | |
| #[test] | |
| fn short_words_spread_across_the_table() { | |
| let cache = WordCache::new(1 << 12); | |
| let homes: std::collections::HashSet<usize> = (0..1000) | |
| .map(|i| { | |
| let (_, hash) = make_word_key(format!("tok{i}").as_bytes()); | |
| hash as usize & cache.index_mask | |
| }) | |
| .collect(); | |
| // 1000 words over 4096 slots share homes by chance alone; ~887 distinct is | |
| // as good as a perfect hash gets, so the floor is well under it. | |
| assert!(homes.len() > 820, "only {} distinct homes", homes.len()); | |
| } | |
| /// Both sides of the 15-byte packing boundary and both sides of the inline/ | |
| /// spill boundary have to survive a round trip, including the long word whose stored | |
| /// key is nothing but a hash. | |
| #[test] | |
| fn every_slot_shape_round_trips() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let long = vec![b'x'; 200]; | |
| let cases: [(&[u8], Vec<u32>); 6] = [ | |
| (b"short", vec![1]), | |
| (b"short-wide", (0..40).collect()), | |
| (b"fifteen-bytes.", vec![2]), | |
| (b"sixteen-bytes.aa", vec![3]), | |
| (&long, vec![9]), | |
| (&long[..64], (0..64).collect()), | |
| ]; | |
| for (word, ids) in &cases { | |
| store(&mut cache, word, ids.clone().into_iter()); | |
| } | |
| for (word, ids) in &cases { | |
| assert_eq!(cache.lookup(word).hit(), Some(&ids[..]), "{word:?}"); | |
| } | |
| } | |
| /// A word past fifteen bytes is stored as a hash, so its length stops mattering and | |
| /// there is no size the cache turns away. | |
| #[test] | |
| fn a_very_long_word_round_trips() { | |
| let mut cache = WordCache::new(1 << 8); | |
| let word = vec![b'q'; 8192]; | |
| let ids: Vec<u32> = (0..2000).collect(); | |
| store(&mut cache, &word, ids.clone().into_iter()); | |
| assert_eq!(cache.lookup(&word).hit(), Some(&ids[..])); | |
| } | |
| /// A tag is one byte, so one slot in 255 answers for a word that is not in | |
| /// it. That is too rare to wait for, so forge one: put a word's tag on a slot | |
| /// holding a different key, and demand the walk look past it rather than | |
| /// treat the tag as the answer. | |
| #[test] | |
| fn a_tag_collision_is_confirmed_against_the_key() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let (_, hash) = make_word_key(b"beta"); | |
| let decoy = hash as usize & cache.index_mask; | |
| cache.tags[decoy] = make_tag(hash); | |
| // Any key that is not beta's. A packed key always carries a length in its | |
| // top byte, so 1 cannot be one. | |
| cache.cached_words[decoy] = CachedWord { | |
| word_bytes_or_hash: 1, | |
| word_ids: [7, 0, 0], | |
| inline_id_count: 1, | |
| }; | |
| store(&mut cache, b"beta", [2u32].into_iter()); | |
| assert_eq!(cache.lookup(b"beta").hit(), Some(&[2u32][..])); | |
| } | |
| /// A word that hashes into a full window is stored rather than turned away, and | |
| /// nothing else in the window is disturbed by it. | |
| #[test] | |
| fn a_full_window_evicts_its_home_slot() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let words: Vec<Vec<u8>> = (0..WALK_WINDOW as u8).map(|i| vec![i; 4]).collect(); | |
| for (i, word) in words.iter().enumerate() { | |
| store(&mut cache, word, [i as u32].into_iter()); | |
| } | |
| let (_, hash) = make_word_key(b"newcomer"); | |
| let home = hash as usize & cache.index_mask; | |
| let evicted = words | |
| .iter() | |
| .position(|word| slot_of(&cache, word) == Some(home)) | |
| .expect("the table is full, so some word holds that slot"); | |
| store(&mut cache, b"newcomer", [999u32].into_iter()); | |
| assert_eq!(cache.lookup(b"newcomer").hit(), Some(&[999u32][..])); | |
| assert_eq!(cache.lookup(&words[evicted]).hit(), None); | |
| for (i, word) in words.iter().enumerate() { | |
| if i != evicted { | |
| assert_eq!( | |
| cache.lookup(word).hit(), | |
| Some(&[i as u32][..]), | |
| "entry {i} was dropped as well" | |
| ); | |
| } | |
| } | |
| } | |
| /// A window that runs off the end of the table carries on at the start, and a | |
| /// word stored past that seam has to be found there. Nothing mirrors the first | |
| /// tags at the end of the row. Every step of the walk is folded back into the | |
| /// table instead. | |
| #[test] | |
| fn a_word_stored_past_the_end_of_the_table_is_found() { | |
| let mut cache = WordCache::new(WALK_WINDOW); | |
| let last_slot = cache.cached_words.len() - 1; | |
| let homed_on_the_last_slot: Vec<Vec<u8>> = (0..4000u32) | |
| .map(|i| format!("w{i}").into_bytes()) | |
| .filter(|word| { | |
| let (_, hash) = make_word_key(word); | |
| hash as usize & cache.index_mask == last_slot | |
| }) | |
| .take(3) | |
| .collect(); | |
| assert_eq!(homed_on_the_last_slot.len(), 3); | |
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | |
| store(&mut cache, word, [i as u32].into_iter()); | |
| } | |
| for (i, word) in homed_on_the_last_slot.iter().enumerate() { | |
| assert_eq!(cache.lookup(word).hit(), Some(&[i as u32][..]), "{word:?}"); | |
| } | |
| // The first one took the last slot, so the other two could only go past it. | |
| assert!(slot_of(&cache, &homed_on_the_last_slot[1]).unwrap() < last_slot); | |
| } | |
| /// Churn a table far too small for the input and demand the invariant that matters: | |
| /// an entry may be evicted, but a hit is never wrong. | |
| #[test] | |
| fn a_hit_never_returns_another_words_ids() { | |
| let mut cache = WordCache::new(64); | |
| let mut expected: Vec<(Vec<u8>, Vec<u32>)> = Vec::new(); | |
| for i in 0..2000usize { | |
| let word = match i % 4 { | |
| 0 => format!("w{i}"), | |
| 1 => format!("a-long-word-past-fifteen-bytes-{i}"), | |
| 2 => format!("k{i}xxxxxxxxxxxx"), | |
| _ => format!("{}-{i}", "z".repeat(i % 40)), | |
| } | |
| .into_bytes(); | |
| let ids: Vec<u32> = (0..=(i % 9) as u32).map(|k| i as u32 * 16 + k).collect(); | |
| store(&mut cache, &word, ids.clone().into_iter()); | |
| expected.push((word, ids)); | |
| } | |
| let mut live = 0; | |
| for (word, ids) in &expected { | |
| if let Some(hit) = cache.lookup(word).hit() { | |
| assert_eq!(hit, &ids[..], "{word:?}"); | |
| live += 1; | |
| } | |
| } | |
| assert!( | |
| live > 0, | |
| "everything was evicted, so the test proves nothing" | |
| ); | |
| } | |
| /// Without the flush the ids buffer grows until its budget is spent and the cache then | |
| /// silently stops accepting any word with more than [`MAX_INLINE_IDS`] ids, even though | |
| /// the table has room. Store far past the budget and demand the last word still landed. | |
| #[test] | |
| fn a_spent_ids_buffer_empties_the_table_instead_of_wedging_it() { | |
| let mut cache = WordCache::new(64); | |
| let word = |i: usize| format!("a-long-word-past-fifteen-bytes-{i:04}"); | |
| for i in 0..5000usize { | |
| store(&mut cache, word(i).as_bytes(), [i as u32; 8].into_iter()); | |
| } | |
| assert!(cache.spilled_ids.len() <= cache.spilled_ids_budget); | |
| store(&mut cache, word(5000).as_bytes(), [5000u32; 8].into_iter()); | |
| assert_eq!( | |
| cache.lookup(word(5000).as_bytes()).hit(), | |
| Some(&[5000u32; 8][..]), | |
| "inserts stopped landing: the buffer never came back" | |
| ); | |
| } | |
| } | |
| output.extend(ids.iter().map(|&id| PipelineToken { id })); | ||
| return Ok(()); | ||
| } | ||
| Lookup::Miss(at) => placement = at, |
Collaborator
There was a problem hiding this comment.
Miss carries the Placement itself now.
Suggested change
| Lookup::Miss(at) => placement = at, | |
| Lookup::Miss(at) => placement = Some(at), |
Comment on lines
+449
to
+451
| if let Some(at) = placement { | ||
| word_cache.insert(at, output[start..].iter().map(|token| token.id)); | ||
| } |
Collaborator
There was a problem hiding this comment.
placement is a Placement, not an Option.
Suggested change
| if let Some(at) = placement { | |
| word_cache.insert(at, output[start..].iter().map(|token| token.id)); | |
| } | |
| word_cache.insert(placement, output[start..].iter().map(|token| token.id)); |
| output.extend(ids.iter().map(|&id| PipelineToken { id })); | ||
| return Ok(()); | ||
| } | ||
| Lookup::Miss(at) => placement = at, |
Collaborator
There was a problem hiding this comment.
Same here -- the local Option still carries "sampling skipped the lookup".
Suggested change
| Lookup::Miss(at) => placement = at, | |
| Lookup::Miss(at) => placement = Some(at), |
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.
TL;DR
Implement a cache to memoize pretoken -> encoded ids
The cache is a lookup table with linear probing, capped at 16 steps
When no empty slot is found during the linear probing, evict naively the original index
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 sweepcc76dfcde · 2026-07-30 17:03 UTC· Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · 16 coresvs base branch (
28b2b9633) — 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 · ×5.14 vs v0.23.1 · ×1.02 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 11+0 (peak 12) · Pipeline 8+2 (peak 17)
deepseek-v4 — deepseek 3-regex split-heavy byte-level BPE · ×23.25 vs v0.23.1 · ×5.20 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) · ×2.18 vs v0.23.1 · ×1.11 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 · ×26.13 vs v0.23.1 · ×3.13 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) · ×21.49 vs v0.23.1 · ×4.11 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) · ×23.12 vs v0.23.1 · ×3.20 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.40 vs v0.23.1 · ×1.13 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 19+0 (peak 23) · Pipeline 23+0 (peak 23)
llama-3 — cl100k-regex byte-level BPE (llama-3), single regex · ×22.76 vs v0.23.1 · ×3.21 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 73+0 (peak 95) · Pipeline 93+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) · ×8.12 vs v0.23.1 · ×1.45 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 152+0 (peak 194) · 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