Skip to content

perf(pre_tokenizers): eliminate per-token heap allocations in ByteLevel decoder - #2376

Open
bsachart wants to merge 1 commit into
huggingface:mainfrom
bsachart:perf/optimize-byte-level-decoder
Open

perf(pre_tokenizers): eliminate per-token heap allocations in ByteLevel decoder#2376
bsachart wants to merge 1 commit into
huggingface:mainfrom
bsachart:perf/optimize-byte-level-decoder

Conversation

@bsachart

@bsachart bsachart commented Sep 1, 2026

Copy link
Copy Markdown

ByteLevel::decode_chain previously allocated a temporary Vec<u8> per token via try_fold(vec![], ...), incurring $N$ dynamic heap allocations per decode call.

This PR replaces per-token allocations with a single pre-allocated buffer sized to tokens.iter().map(String::len).sum(). Tokens containing special/unknown characters trigger an $O(1)$ bytes.truncate(start) rollback before inserting raw bytes, eliminating intermediate vector clones. CHAR_BYTES is derived dynamically from BYTES_CHAR in one line.

Benchmark Results (cargo bench --bench decode_benchmark -- '^decode-llama3-(en|ja)/decode$')

  • Llama 3 English decode: 157.82 ms → 95.42 ms (~39.5% faster)
  • Llama 3 Japanese decode: 25.22 ms → 18.83 ms (~25.3% faster)

All 204 library unit tests pass.

AI-assisted change.

… 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 ArthurZucker left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ty, tho we should be comparing againt this: #2119

@bsachart

bsachart commented Sep 2, 2026

Copy link
Copy Markdown
Author

Thanks for the pointer! I checked out #2119 (feat/train_encode_split) and saw that ByteLevelDecoder in tk-encode/src/decoders/byte_level.rs still has the per-token try_fold(vec![], ...) allocation loop.

I tested the single-pass pre-allocated buffer fix directly on ByteLevelDecoder::decode_chain in #2119. On 5M tokens, execution time dropped from ~334ms to ~220ms (~34% faster latency / ~1.5x throughput).

Would you prefer I retarget this PR to feat/train_encode_split, or open a new PR against #2119?

@ArthurZucker

Copy link
Copy Markdown
Collaborator

Yes please target the new branch!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants