perf(cache): take a word key in one load instead of a memcpy - #2276
Open
ArthurZucker wants to merge 1 commit into
Open
perf(cache): take a word key in one load instead of a memcpy#2276ArthurZucker wants to merge 1 commit into
ArthurZucker wants to merge 1 commit into
Conversation
pack_word copied `len` bytes at a time, which is a call into memcpy and a branch that mispredicts on every short word. Read the 16 bytes at once and mask the surplus, declining the wide load when it would cross a page.
|
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. |
ArthurZucker
commented
Aug 4, 2026
| let mut lanes = [0u8; 16]; | ||
| lanes[..len].copy_from_slice(word); | ||
| Some(u128::from_le_bytes(lanes) | ((len as u128) << 120)) | ||
| let ptr = word.as_ptr(); |
Collaborator
Author
There was a problem hiding this comment.
we are masking raw even in the else branch which is probably not optimal 👀
ff to push to fix if this has an impact!
This was referenced Aug 4, 2026
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
2 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.
Closes the
TODOonpack_word.pack_wordbuilt the key withlanes[..len].copy_from_slice(word).lenis onlyknown at run time, so that is a call into
memcpybehind a length branch thatmispredicts on nearly every short word. Read the 16 bytes at
word.as_ptr()in oneunaligned load and mask the surplus off; the load reads past the word, so it is only
taken when all 16 bytes fall inside the page the word already sits in, and the old
copy stays as the fallback for the rest.
The packed
u128is byte-identical to before — same value, same length in the topbyte — so slot placement, tags and key comparison are unchanged. A new test asserts
both reads agree, including a word deliberately placed to straddle a page boundary.
Measured
Key construction timed on its own (spans from a pre-tokenizer pass over each corpus,
then key + hash per word), aarch64,
lto = "fat". ns per word, so it is independentof how the text was split:
copy_from_sliceFor scale: on english the old key path cost 1.59 ns/B against 0.48 ns/B for the whole
pre-tokenizer split, i.e. building the key cost more than 3× cutting the text into
words. It is now 0.36 ns/B.
Two variants I measured and did not ship:
crc32cdinstead of aHash on the packedu128— a further 1.7× on english(0.36 → 0.21 ns/B), but nothing on russian, and it needs a
cfgfallback plusre-seeding to keep the table randomised. Happy to add it as a second commit if you
want it.
reads covering
[0, len)exactly, no reading past the word). Byte-identical andneeds no
unsafe, but the 3-way length branch costs most of the win back:0.92 ns/B vs 0.36 on english.
vs #2267
Orthogonal, and they compose. #2267 changes where keys are computed (keyed spans
batched ahead of the model); this changes what one key costs, and does not touch the
pipeline. #2267 leaves
pack_wordas it was, so it inherits the same 4.6× if both land.Worth noting for #2267's own question: in an A/B on this machine, computing the key
inside the split's emit walk was 7–22% slower than a separate pass over the emitted
spans on every corpus — the fused loop serialises
open → key → next, where two loopslet the span array prefetch. That is the key-placement question only; it says nothing
about batching into
tokenize_keyed_spans, which #2267's own harness measures (I couldnot run it here,
data/fixtures/is empty locally).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 sweepe38b966bc · 2026-08-04 04:18 UTC· Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz · 48 coresvs base branch (
e87a72c56) — 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.23 vs v0.23.1 · ×1.06 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 12+0 (peak 12) · Pipeline 8+2 (peak 17)
deepseek-v4 — deepseek 3-regex split-heavy byte-level BPE · ×25.28 vs v0.23.1 · ×6.41 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 62+0 (peak 67) · 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.10 vs v0.23.1 · ×1.17 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 · ×30.08 vs v0.23.1 · ×3.82 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 25+2 (peak 27) · Pipeline 27+0 (peak 27)
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) · ×23.22 vs v0.23.1 · ×5.12 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.81 vs v0.23.1 · ×3.91 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 169+0 (peak 230) · 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.21 vs v0.23.1 · ×1.15 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 · ×25.81 vs v0.23.1 · ×3.85 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 73+0 (peak 95) · Pipeline 92+0 (peak 94)
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) · ×7.94 vs v0.23.1 · ×1.55 vs base · decode pending
Memory (RSS MB, load+encode): v0.23.1 152+0 (peak 194) · Pipeline 109+0 (peak 194)
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