feat(bpe): ScratchPool pattern + BPE word cache#2223
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. |
9582b69 to
1f9f034
Compare
SBrandeis
changed the title
| /// naive implem of word -> IDs cache | ||
| pub struct WordCache { | ||
| capacity: usize, | ||
| lookup: AHashMap<String, Box<[u32]>>, | ||
| } |
There was a problem hiding this comment.
YES (look at the perf improvements)
| let maybe_scratch = self | ||
| .pool | ||
| .lock() | ||
| .unwrap_or_else(PoisonError::into_inner) |
There was a problem hiding this comment.
what does this do practically? does it return Vec<PipelineModelScratch> in its last known state?
There was a problem hiding this comment.
PoisonError::into_inner ignores the poisoning and hands the lock regardless of the error, which is safe here because what we do is push/pop on a Vec
Worst case, a scratch buffer is lost or leaked, which is fine because we clear it at L303
There was a problem hiding this comment.
From Claude:
1. The only code that can panic while holding this lock is Vec::push's capacity overflow — practically unreachable here (pop is panic-free, OOM aborts).
2. Even in that unreachable case, the Vec is left valid and merely unchanged, and get already tolerates a smaller-than-expected pool via model.init_scratch() at line 306.
| .pool | ||
| .lock() | ||
| .unwrap_or_else(PoisonError::into_inner) | ||
| .pop(); // Release the lock |
There was a problem hiding this comment.
we sure lock is released here?
There was a problem hiding this comment.
used a {} block to make it more obvious
| impl ScratchPool { | ||
| fn new() -> Self { | ||
| Self { | ||
| pool: Mutex::new(Vec::new()), |
There was a problem hiding this comment.
shouldn't it be Arc<Mutex<...>> instead since we'll be sharing this across threads?
| // Option so we can use `.take()` to reclaim ownership of the scratch | ||
| // to push it back to the pool | ||
| scratch: Option<PipelineModelScratch>, | ||
| pool: &'a ScratchPool, |
There was a problem hiding this comment.
| pool: &'a ScratchPool, | |
| scratch_pool: &'a ScratchPool, |
| self.pool | ||
| .pool | ||
| .lock() | ||
| .unwrap_or_else(PoisonError::into_inner) |
There was a problem hiding this comment.
I understand that this allows some form of recovery. If a thread panics, do we want to abort the program or recover the last known state of the pool?
| impl std::ops::Deref for ScratchGuard<'_> { | ||
| type Target = PipelineModelScratch; | ||
| fn deref(&self) -> &PipelineModelScratch { | ||
| self.scratch.as_ref().unwrap() |
There was a problem hiding this comment.
changed the type to remove the necessity of unwrap
|
|
||
| impl std::ops::DerefMut for ScratchGuard<'_> { | ||
| fn deref_mut(&mut self) -> &mut PipelineModelScratch { | ||
| self.scratch.as_mut().unwrap() |
|
|
||
| /// The default capacity for a `BPE`'s internal cache. | ||
| pub static DEFAULT_CACHE_CAPACITY: usize = 10_000; | ||
| pub static DEFAULT_CACHE_CAPACITY: usize = 1 << 16; |
There was a problem hiding this comment.
bitshift + usize seems error prone, although shifting by 16 should be safe for most platforms
SBrandeis
changed the title
9253dfa to
1058288
Compare
1058288 to
51d2ab0
Compare
| self.scratch_pool | ||
| .pool | ||
| .lock() | ||
| .unwrap_or_else(PoisonError::into_inner) |
There was a problem hiding this comment.
does acquiring the lock + calling poisonerror::into_inner "unpoison" the lock?
Just found: mutex.clear_poison(); might be worth calling, wonder if we can do it in the unwrap_or_else closure
There was a problem hiding this comment.
does acquiring the lock + calling poisonerror::into_inner "unpoison" the lock?
It does not
Just found: mutex.clear_poison(); might be worth calling, wonder if we can do it in the unwrap_or_else closure
Wouldn't do much except an unecessary flag flip, no?
There was a problem hiding this comment.
Feels weird to continue using a poisoned mutex, but from what I can gather it's not much of an issue indeed. If a thread were to fail a write to the inner pool object, could we leave things in a state that would cause downstream errors (e.g. partial writes) or is it simply a matter of "has been added/removed to the pool" and since it's a cache we don't really care?
There was a problem hiding this comment.
It is mostly safe to fail at that specific boundary for a few reasons:
- the only thing that can panic while holding the lock is the
Vec::push(out of capacity) - even if something else would panic, the lock is only held to push or pop a scratch buffer into or from the pool
- not managing to re-pool a scratch buffer is safe (we lost the cache but not critical)
- not managing to pop an existing scratch buffer is safe, the code defaults to
::init()ing a fresh one
With the current state of the code, the Mutex being poisonned is both unlikely and not recoverable
3 participants
PipelineTokenizer benchmark
7 / 8 models supported — PipelineTokenizer vs
tokenizersv0.23.1 (latest release) · ~10 kB inputs · single thread + 1/2/4/8/max-thread sweep51d2ab06f · 2026-07-22 08:29 UTC· Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · 48 coresvs base branch (
575582072) — per-model geomean ×speedup of this PR's PipelineTokenizer against the base branch's; regressions in red.bert-base-uncased — normalizer-heavy WordPiece · ×4.95 vs v0.23.1 · ×1.01 vs base
Memory (RSS MB, load+encode): v0.23.1 12+0 (peak 12) · Pipeline 8+0 (peak 17)
deepseek-v4 — deepseek 3-regex split-heavy byte-level BPE · ×22.28 vs v0.23.1 · ×5.45 vs base
Memory (RSS MB, load+encode): v0.23.1 62+0 (peak 68) · Pipeline 96+0 (peak 96)
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).gpt2 — gpt2 ByteLevel regex · ×26.95 vs v0.23.1 · ×3.67 vs base
Memory (RSS MB, load+encode): v0.23.1 25+2 (peak 27) · Pipeline 30+0 (peak 29)
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) · ×24.52 vs v0.23.1 · ×2.97 vs base
Memory (RSS MB, load+encode): v0.23.1 4+11 (peak 16) · Pipeline 7+2 (peak 7)
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.89 vs v0.23.1 · ×1.91 vs base
Memory (RSS MB, load+encode): v0.23.1 4+11 (peak 15) · Pipeline 6+2 (peak 7)
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 · ×3.72 vs v0.23.1 · ×1.16 vs base
Memory (RSS MB, load+encode): v0.23.1 19+0 (peak 23) · Pipeline 22+0 (peak 21)
llama-3 — cl100k-regex byte-level BPE (llama-3), single regex · ×21.20 vs v0.23.1 · ×3.20 vs base
Memory (RSS MB, load+encode): v0.23.1 143+0 (peak 199) · Pipeline 144+0 (peak 200)
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