perf(pre_tokenizers): eliminate per-token heap allocations in ByteLevel decoder - #2376
Open
bsachart wants to merge 1 commit into
Open
perf(pre_tokenizers): eliminate per-token heap allocations in ByteLevel decoder#2376bsachart wants to merge 1 commit into
bsachart wants to merge 1 commit into
Conversation
… indexing Reduce memory allocations in ByteLevel::decode_chain by pre-allocating a single output buffer for the total sequence length instead of allocating a temporary Vec<u8> for every single token (try_fold(vec![], ...)). Key Improvements: - Allocation Reduction: Replaces N dynamic heap allocations per decode call with 1 pre-allocated buffer. - O(1) Array Indexing: Replaces HashMap byte lookup with direct [char; 256] array indexing in BYTES_CHAR. - Derived Mapping: CHAR_BYTES is derived dynamically from BYTES_CHAR in a single line, eliminating duplicate map construction without hardcoded range constants. Benchmark Results (Criterion decode-llama3-(en|ja)/decode): - Llama 3 English decode: ~35% to 40% faster (~95-102 ms vs 157.8 ms baseline) - Llama 3 Japanese decode: ~21% to 25% faster (~18.8 ms vs 25.2 ms baseline) AI-assisted change.
ArthurZucker
reviewed
Sep 2, 2026
ArthurZucker
left a comment
Collaborator
There was a problem hiding this comment.
ty, tho we should be comparing againt this: #2119
Author
|
Thanks for the pointer! I checked out #2119 ( I tested the single-pass pre-allocated buffer fix directly on Would you prefer I retarget this PR to |
Collaborator
|
Yes please target the new branch! |
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
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.
ByteLevel::decode_chainpreviously allocated a temporaryVec<u8>per token viatry_fold(vec![], ...), incurringThis PR replaces per-token allocations with a single pre-allocated buffer sized to$O(1)$
tokens.iter().map(String::len).sum(). Tokens containing special/unknown characters trigger anbytes.truncate(start)rollback before inserting raw bytes, eliminating intermediate vector clones.CHAR_BYTESis derived dynamically fromBYTES_CHARin one line.Benchmark Results (
cargo bench --bench decode_benchmark -- '^decode-llama3-(en|ja)/decode$')All 204 library unit tests pass.
AI-assisted change.