SIGN IN SIGN UP

perf(toktrie): Optimize Trie initialization to eliminate allocation storms (#344)

This commit refactors the initialization and filtering of `TokTrie` to
bypass severe allocator lock-contention, particularly when running in
multi-threaded, security-hardened environments like Chromium
(PartitionAlloc).

**The Problem:**

Previously, building the temporary `TrieHash` utilized nested vectors
(`children: Vec<TrieHash>`). For a 200k plus token vocabulary, this
triggered ~860,000 individual `Vec::push` reallocation calls during
construction, and ~860,000 `drop_in_place` / `madvise` deallocation
calls during cleanup. In environments like Chromium, this massive
allocation churn causes severe thread-cache exhaustion and global lock
contention, resulting in 10-second initialization hangs on background
threads.

**The Solution:**
1. **Flat Arena (`TrieBuilder`):** Replaced `TrieHash` with an
arena-allocated `Vec<BuilderNode>`, where nodes reference children and
siblings via `u32` indices. This builds the entire tree using exactly
one system allocation and instantly drops it with one `free`.

2. **O(1) Insertion Paths:** Added `root_children` and `last_child`
pointers to the builder. This turns O(N) sibling scans into strict O(1)
appends, while carefully preserving the duplicate-token ID insertion
order relied upon by the Earley parser.

3. **Allocation-Free Serialization:** By guaranteeing tokens are
inserted alphabetically, siblings are naturally sorted in the arena.
`serialize_node` now traverses the linked list directly without
allocating temporary `Vec`s to sort children.

4. **Sub-Second Filtering (`sorted_vocab`):** Cached the alphabetically
sorted token indices on `TokTrie`. `TokTrie::filter` now iterates this
cache directly, bypassing hundreds of thousands of `Vec<u8>` string
allocations and sorting overhead during active constrained decoding.

**Results:**
- Eliminates >1.7 million `malloc` / `free` operations per tokenizer
initialization.
- Reduces raw initialization time in isolated Release builds from ~3.0s
to ~0.8s (>60% CPU-bound speedup).

---------

Co-authored-by: Nathan Memmott <memmott@chromium.org>
N
nathanmemmott committed
97980c5059f60e591c2e678a9b963c6bd11b4e41
Parent: b01742e
Committed by GitHub <noreply@github.com> on 4/28/2026, 11:58:15 PM