perf(decode): index the vocab slab by id, not through the MPHF slot - #2387
Open
ArthurZucker wants to merge 9 commits into
Open
perf(decode): index the vocab slab by id, not through the MPHF slot#2387ArthurZucker wants to merge 9 commits into
ArthurZucker wants to merge 9 commits into
Conversation
`decode_byte_level` reached every token through `id_to_token_bytes`, which is a three-load pointer chase -- `id_to_slot[id]` -> `spans[slot]` -> `bytes` -- and whose middle hop is a *random* access, because slot order is the MPHF's and not id order. `cargo asm` on the loop shows that chase plus three bounds checks per token, against a flat-arena decoder's one load and one copy. `build` now lays the slab out in ascending id order (free: nothing below the placement loop depends on iteration order -- `entries`, `spans` and `id_to_slot` are written by slot or by id, and `keys` only feeds an order-independent `HashSet`) and records `decode_off`, so a decoder reads an adjacent `u32` pair, usually from one cache line, and slices. An id the vocabulary does not hold gets equal offsets, i.e. an empty slice, which is what a decoder appends for it anyway. Costs 4 bytes per id in the id space, ~513 kB on a 128k-id vocabulary, and duplicates no token bytes. Measured, llama-3, `tk-serialize`'s decode bench in 10 kB chunks (p = 0.00): english 367.4 -> 428.0 MiB/s +16.5% japanese 449.2 -> 490.1 MiB/s +9.1% And in huggingface/tokbench, where every engine's decoded text is hashed and compared against the released crate -- all cells byte-exact: corpus before after tokie was now english 380.8 464.4 466.3 0.841x 0.996x chinese 353.3 401.3 386.0 0.924x 1.04x japanese 390.0 442.6 425.2 0.954x 1.04x Two things that did NOT work, recorded so they are not retried: raising the output `Vec` estimate from 4 to 5 bytes per token (english flat, japanese -4.9% -- over-reserving costs more than the grow it avoids), and outlining the added-token arm as `#[inline(never)] #[cold]` (freed the spilled base pointers into registers and cut the loop from 420 to 348 instructions, and moved the clock not at all -- those spills were L1 hits). The bench that found this is new. The existing decode benches call `decode` once per *line*, where one allocation and one `from_utf8` validation per short line dominate and nothing in the token loop is visible; the 10 kB chunk bench matches the regime tokbench measures.
Two phases instead of one loop: phase 1 walks the ids collecting each token's byte slice, phase 2 concatenates them into an exactly-sized buffer. The one-pass loop interleaved a random probe into the decode index with an append to the output, so every token's copy waited behind its own lookup. Separated, phase 1 issues nothing but independent random loads with no output cursor to serialise on, and phase 2 is a sequential walk. Phase 1 also yields the exact output length, retiring the `ids.len() * 4` guess and the realloc it caused on Latin text. llama-3, 10 kB chunks, criterion: english 428 -> 501 MiB/s +17.2% japanese 490 -> 504 MiB/s +2.8% In huggingface/tokbench, every engine's decoded text hashed and compared against the released crate -- all cells byte-exact: corpus before after tokie english 464.4 533.1 456.3 now 1.17x ahead japanese 442.6 452.5 423.3 chinese 401.3 395.1 386.4 Three variants of this lost, recorded in the doc comment so they are not retried: 4- and 16-lane splits of the id stream with per-lane output buffers (-16% and -21%, overhead scaling with lane count -- the parallelism was already being extracted by the out-of-order engine, the bookkeeping was not free); parking `(start, end)` in a reused thread-local scratch rather than `&[u8]` in a fresh `Vec` (-3.3%, because phase 2 then re-slices the slab with a bounds check per token); and prefetching the index eight tokens ahead (+3.1% english, flat japanese, not worth arch-gated inline asm). Sequences containing an added token keep the one-pass walk: those yield an owned `String` rather than a borrow into the vocabulary, so they cannot be gathered as slices.
Two changes to the two-phase path, measured separately. Phase 1 already computes the exact output length, so `extend_from_slice`'s per-token capacity check and length store can only ever pass: write straight through the pointer and `set_len` once. English +1.8%, japanese +0.8%. Phase 1's gather buffer becomes thread-local, allocated on a thread's first decode and reused after, retiring a malloc plus a first touch of every page of it on every call -- on a ~2250-token chunk producing 10 kB of text, a 36 kB allocation written and then read straight back to move 10 kB. English +1.3%, japanese +0.5%. The payload stays `(ptr, len)` rather than `(start, end)` into the slab, and that is the whole reason this works now: an earlier attempt paired buffer reuse with offsets and came out 3.3% DOWN, because phase 2 then re-slices the slab with a bounds check per token. Reuse was worth +1.3% all along; the offsets were costing more than the malloc. One change at a time. Cumulative, llama-3, 10 kB chunks, criterion: english 501 -> 517 MiB/s japanese 504 -> 511 MiB/s And in huggingface/tokbench, all cells byte-exact against the released crate: corpus decode tokie english 537.3 465.3 1.15x japanese 457.3 396.6 1.15x chinese 419.1 376.0 1.11x `decode` still allocates its returned `String`, which is the one allocation left in the path.
Decode's cost is one random probe per token into an index sized by the id space, so vocabulary size is a first-order variable that a single model cannot show. Three configs, one bench: llama-3 ByteLevel, 128k ids -> ~513 kB index gpt-oss ByteLevel, 200k ids -> ~800 kB index gemma BPE + byte_fallback, Sequence[Replace, ByteFallback, Fuse] The first two answer the residency question: 511 and 514 MiB/s on english, statistically identical, so a 1.56x larger index costs nothing at these sizes. That also retires the idea that prefetching the index would pay on a bigger vocabulary -- it does not, because the index was never the miss it looked like. The third is why this is worth committing. gemma measures **34 MiB/s against 514**, ~15x, because it is not byte-level: `decode` falls through to a route that calls `id_to_token` per id, which allocates a `String` every time, collects a `Vec<String>`, and only then runs the decoder chain -- 1.5M allocations for one pass over big.txt. Every SPM-family config is on that path (gemma, llama-2, mistral, T5), and no bench in the tree showed it. Left as a standing measurement rather than fixed here: precomputing the chain per vocabulary entry at load, the way `byte_level::transform_vocab` already does, gets gemma to 349 MiB/s but is NOT correct. `ByteFallback` is run-based rather than per-token, and beyond that SPM's leading `▁` makes the decoder's prefix handling position-dependent, which a per-token table cannot express. Verified against the released crate: fixing the byte-fallback half repairs gemma's CJK cells, and gemma/english plus all three llama-2 cells still differ.
Every SentencePiece-family config -- gemma, llama-2, mistral, T5 -- decoded
through a route that allocates roughly three `String`s per token: one in the
gather (`id_to_token` is `from_utf8_lossy(..).into_owned()`), one in `Replace`
(`"".to_string()` per token), a fresh `Vec<String>` plus a
`previous_byte_tokens.clone()` per byte run in `ByteFallback`, then `join("")`
in `Fuse`. About 4.5M allocations for one pass over big.txt, and 34 MiB/s
against 514 for a byte-level model on the same corpus.
`Sequence[Replace{literal}, ByteFallback, Fuse]`, optionally followed by
`Strip`, is now recognised at load and run as a single pass into a reused
buffer: `Replace` inline per token, `ByteFallback` accumulating a run and
resolving it at the end, `Fuse` being the concatenation the loop already does,
`Strip` applied to the finished output. Nothing is allocated per token. An
unrecognised chain declines and keeps the generic route.
Measured, verified byte-exact against the released crate in
huggingface/tokbench (`decode_verified` on all 9 cells, encode ids too):
gemma-3 chinese 39.4 -> 190.6 MB/s 4.8x
english 37.2 -> 177.9 4.8x
japanese 44.1 -> 218.4 5.0x
llama-2 now 198-273 MB/s on the same path
llama-3 unchanged (byte-level, 401-553)
And on the larger gemma-4 config in `tk-serialize`'s bench: english
34.4 -> 197.8 MiB/s, japanese 45.1 -> 303.4.
Two things this gets right that a cheaper version did not, both caught by that
gate rather than by reasoning:
* `ByteFallback` is run-based, not per-token. Precomputing the chain per
vocabulary entry asks it to validate one byte alone; a lone `0xE4` is not
UTF-8 and comes back U+FFFD, so every three-token CJK character became
three replacement characters. An **invalid** run must yield one U+FFFD per
byte of the run, which "append raw bytes and let `from_utf8_lossy` sort it
out" also does not reproduce.
* `Strip` runs *after* `Fuse`, so it applies to the whole output, not to each
token. That position-dependence is why a per-token table cannot express
this chain, and it broke every llama-2 cell before being handled here.
Membership in the added vocabulary is a binary search over the added ids rather
than `id >= added_id_min`: that short-circuit never fires for these configs,
because gemma's `<pad>` is id 0 and so every token clears the threshold. Worth
knowing for #2178, which proposes exactly that short-circuit.
The fused SPM pass still ran `Replace` against every token *occurrence*. A
byte-level model does not: `byte_level::transform_vocab` decodes its whole
vocabulary at load, which is most of why that path is a gather and a `memcpy`
and this one was not. Same treatment here -- transform once per vocabulary
*entry*, at load.
A byte-fallback entry cannot be pre-transformed, since it has to rejoin a run
at decode time, and its length does not identify it -- plenty of ordinary
tokens are also one byte. The flag therefore lives in **bit 31 of the start
offset**, so the adjacent-pair load a decoder already performs yields the span
and the flag together: no second array, no second cache line, no extra probe.
Slabs are far below 2 GB, so 31 bits of offset costs nothing.
The decode loop's ordinary arm is now `extend_from_slice` and nothing else.
Measured, `tk-serialize`'s bench on gemma-4, 10 kB chunks:
english 206.7 -> 310.8 MiB/s +50% (34.4 at the start of this work, 9.0x)
japanese 290.3 -> 383.3 +32% (45.1 at the start, 8.5x)
And in huggingface/tokbench, byte-exact against the released crate on all nine
cells (`decode_verified` and encode `verified` both true):
gemma-3 chinese 187.8 -> 214.6 MB/s (39.4 baseline, 5.4x)
english 178.2 -> 229.0 (37.2 baseline, 6.2x)
japanese 213.8 -> 254.6 (44.1 baseline, 5.8x)
llama-2 chinese 208.7 -> 262.3
english 188.0 -> 328.3
japanese 260.3 -> 339.3
llama-3 unchanged (418, 562, 476)
Costs the transformed vocabulary: 4 bytes per id in the id space plus a copy of
the token bytes, so roughly 3 MB on gemma's 262k ids. That is the same trade
`transform_vocab` already makes for byte-level models.
Also measured and NOT taken, on the way here: giving the fused pass an owned
output `Vec` so `String::from_utf8` could take ownership instead of
`from_utf8_lossy(..).into_owned()` copying a reused buffer. The copy is real,
but the malloc that replaces it costs the same -- english +4.5%, japanese
-4.3%. It is a wash here and not for the byte-level path because that one knows
its exact output length from its gather phase, while this loop cannot: `Replace`
changes byte counts and byte-fallback runs collapse.
Profiling said the SPM loop was 12.59 ns/token against a byte-level model's
~8.7, and the token mix said why it was NOT what I had assumed: byte-fallback
tokens are 0% of english and 1.6% of japanese, added tokens 0% of both, so the
run machinery was never the cost. The cost was `is_added`, a binary search over
the added ids, run on *every* token -- 1.53M times for one pass over big.txt,
always returning false.
So the flags move into the high bits of the start offset the decoder already
loads: bit 31 byte-fallback, bit 30 special. One adjacent-pair load per token
now yields the span and every flag, and the ordinary arm is `extend_from_slice`
and nothing else. No binary search, no second array, no second probe, and
`simple_id_to_token` is gone from the loop entirely -- added tokens are ordinary
entries in the table, with their own decoded bytes.
`SPECIAL` is set from `is_special_token`, which tests the token *string*, so a
model-vocabulary entry that happens to spell a special token is treated exactly
as the generic route treated it.
Offsets keep the low 29 bits, capping a slab at 512 MB and asserted at build.
Vocabulary slabs are three orders of magnitude below that.
Measured, byte-exact against the released crate on all nine tokbench cells:
gemma-3 chinese 214.6 -> 374.9 MB/s (39.4 baseline, 9.5x)
english 229.0 -> 372.8 (37.2 baseline, 10.0x)
japanese 254.6 -> 411.0 (44.1 baseline, 9.3x)
llama-2 chinese 262.3 -> 293.1
english 328.3 -> 345.4
japanese 339.3 -> 397.0
llama-3 unchanged (413, 536, 444)
gemma now decodes within 10-30% of a byte-level model on the same corpora,
against 15x slower when this started.
`tk-serialize`'s bench on the larger gemma-4 config agrees: english
310.8 -> 347.4 MiB/s. Note one criterion run reported 149.9 for that cell with
a 121-179 confidence interval; it was contended, and a clean re-run gave 347.4.
A wide interval is the tell.
Also drops the fallback arm that transformed per occurrence. It was unreachable
-- the dispatch already requires BPE and the table is built for every BPE model
-- and an unreachable second implementation of these semantics is exactly the
thing that rots. `recognise` now builds the table itself and declines when it
cannot, so there is one path.
`FusedSpmDecoder` kept `from`/`to` -- the `Replace` pattern -- for the life of
the tokenizer, but the only reader was `push_replaced`, and the only caller of
that is `SpmDecodeTable::build`. Construction-only state does not belong on a
struct that outlives construction.
`push_replaced` becomes a free function taking the pattern, the table's `build`
takes it directly, and the decoder is constructed in one expression instead of
being built half-empty and patched. That also retires the `Default` impl the
two-step construction needed.
`FusedSpmDecoder` is now `{ strip, table }`.
No behaviour change: tests green, clippy and fmt clean, and tokbench reports
zero pipeline mismatches across all nine cells.
Doc comments had accumulated measurement logs and lists of approaches that lost. That belongs in commit messages, which have it; in the source it is prose that goes stale the next time anyone touches the loop. Kept: the bit layout and the 29-bit offset cap, why a byte-fallback entry cannot be pre-transformed, why the slab is in id order, why the added-token path falls back, and every SAFETY comment. Dropped: throughput figures, instruction counts, allocation counts, percentages, and the rejected-variant lists. Also fixes a real defect from an earlier edit of mine: the two-phase explanation had ended up stacked on top of `decode_fused_spm`'s own doc comment, leaving `decode_byte_level` -- the function it describes -- with none. No code change. Tests green, clippy and fmt clean.
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
1 participant
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 decode gap against
tokie, which was the fastest decoder in tokbench and 1.19× ahead of us on english.What was slow
decode_byte_levelreached every token throughid_to_token_bytes:Three dependent loads, and the middle hop is a random access — slot order is the MPHF's, not id order.
cargo asmon the loop shows that chase plus three bounds checks per token, and the vocabulary base pointers spilled to the stack and reloaded every iteration. A flat-arena decoder does one load of an adjacent offset pair and a copy.What changed
buildnow lays the slab out in ascending id order and recordsdecode_off, sobytes[decode_off[id] .. decode_off[id + 1]]is the token. The two offsets are adjacent, usually in the same cache line.The reorder is free: nothing below the placement loop depends on iteration order —
entries,spansandid_to_slotare written by slot or by id, andkeysonly feeds an order-independentHashSet. An id the vocabulary does not hold gets equal offsets, i.e. an empty slice, which is exactly what a decoder appends for it.Cost: 4 bytes per id in the id space — ~513 kB on a 128k-id vocabulary — and no duplicated token bytes.
Measured
tk-serialize's decode bench, llama-3, 10 kB chunks, criterion, p = 0.00 on both:And in tokbench, which hashes every engine's decoded text and compares it against the released crate — all cells byte-exact:
The bench is new, and that matters
The existing decode benches call
decodeonce per line. big.txt lines are short, so one allocation and one fullfrom_utf8validation per line dominate, and nothing happening in the token loop is visible. Concretely: a change that cut the loop from 420 to 348 instructions moved those benches not at all.The added
decode, 10 kB chunksbench uses tokbench's chunking, so it is both loop-sensitive and directly comparable with that matrix.Two things that did not work
Recorded so nobody burns the time again:
Vecestimate fromids.len() * 4to* 5. English output is 4.55 B/token so the 4× guess does under-allocate — but english came out flat (p = 0.12) and japanese regressed 4.9%. Over-reserving 26% costs more than the grow it avoids.#[inline(never)] #[cold]. It worked structurally — confirmed in asm, the base pointers moved out of the stack into registers and the loop went 420 → 348 instructions, 37 → 30 calls — and the clock did not move. Those spills were L1 hits.Tests
decode_index_matches_the_pointer_chaseasserts the new index agrees with the chase it replaces for every id in the id space, on a deliberately sparse vocabulary. Sparse ids are the real risk: a hole is encoded asdecode_off[id] == decode_off[id + 1], and an off-by-one in that walk would hand a hole its neighbour's bytes — a silently wrong decode, not a crash.cargo test -p tk-encode -p tk-serializeis green (176 + 38 + 2 + 2). clippy andfmt --checkclean. The now-deadPipelineBPE::id_to_token_bytespassthrough is removed; the store's ownid_to_token_bytesstays, still used by serialization andid_to_token.