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