perf(pipeline): partition plan_work's chunks instead of sorting them - #2378
Merged
Merged
Conversation
…nning `parallel::encode` gated on `inputs.size_bytes() < PARALLEL_MIN_BYTES`, then planned, then bailed to `encode_serial` when planning produced fewer than two tasks. For one document between `PARALLEL_MIN_BYTES` and twice it that bail is guaranteed: `plan_sequence` short-circuits any sequence under `2 * PARALLEL_MIN_BYTES` into a single chunk, one chunk makes one task, and one task fails the check. So the allocations were provably discarded, and the condition is knowable from the length alone. No measurable throughput effect, and it is worth recording why rather than implying one. Because `plan_sequence` short-circuits *before* its added-token scan, the only waste in that window is about six small `Vec` allocations, against ~22 us of encoding for a 9 kB document. Measured with single-document encodes at 4 kB / 9 kB / 20 kB (the first and last as controls, since the guard cannot fire there): every cell within +-1.5%, including the 9 kB window itself. Kept because it removes work that cannot ever be used, not because it is faster. The expensive version of this waste is at 16 kB and above, where `plan_sequence` does run the full `SpecialSegmentIterator` pass over the document, finds no special token, emits one chunk and discards it -- a duplicated added-token scan. This guard does not cover that, and it cannot cheaply: whether a sequence splits is only knowable from that scan. The real fix there is the intra-sequence splitting #2331 left as follow-up, which would make the planning useful instead of merely cheap to skip.
`plan_work` sorted the chunk list with `sort_unstable_by_key(Reverse(len))` and then folded it a second time for the per-sequence `chunk_count`. Two things are wrong with that shape: - `SequenceChunk` is 40 bytes, so a comparison sort does O(n log n) moves of 40 bytes each. The container was never the problem -- `chunks` is a `Vec`, access was already O(1) -- what it sorts is. - It runs in the serial prologue, on the calling thread, before a single worker is spawned. It is pure Amdahl fraction, and worst exactly in the many-small-inputs case tasks were introduced for in #2338: one chunk per sequence means n is the batch size, so an 80k-document batch sorts 80k structs before any encoding starts. The grouping it feeds never needs a total order -- it walks the chunks accumulating bytes and cuts a task every `PARALLEL_MIN_BYTES`. Ordering by power-of-two size class is sufficient, and that is a counting sort: histogram, exclusive prefix sums largest-class-first, scatter. O(n), moving 4-byte indices instead of 40-byte structs, with `chunk_count` coming off the same histogram pass rather than a second fold. `tasks` now index `order` rather than `chunks`, so the worker pays one extra `u32` load per chunk -- against at least `PARALLEL_MIN_BYTES` of tokenizing per task. Batched `encode` of a 5 MB corpus, min of 5 passes with the corpus repeated to fill 250 ms per pass, 3 interleaved rounds, separately built binaries: 64 B documents n_docs base MB/s this commit gpt2/chinese 80457 362.7 +17.8% llama-3/chinese 80457 344.4 +11.0% llama-3/code 32691 709.0 +6.8% gpt2/english 81983 668.4 +4.0% gpt2/code 32691 683.9 +2.9% llama-3/english 81983 663.8 +1.5% 1024 B documents n_docs base MB/s this commit llama-3/chinese 5118 600.9 +2.1% gpt2/chinese 5118 846.5 +1.8% others -1.0% .. +0.9% Positive in every cell where n is large and gone by 1 kB documents, which is the shape an O(n log n) -> O(n) change in a serial prefix should have. The noise floor for this setup is about +-4%, measured from a commit that cannot affect the batch path at all, so the two chinese cells are the solid ones and the sub-3% entries are marginal. Token ids are unaffected: chunks are encoded independently and reassembled by `(seq, idx)`, so visiting order cannot change the result. Verified identical across all four builds in every cell. Only the task boundaries move.
Replaces the counting sort from the previous commit with a single partition pass, on review against #2361 which reached the same conclusion independently and expressed it in eight lines. The task grouping never needed the ordering the counting sort provided. Tasks are byte-balanced -- each grown until it holds `PARALLEL_MIN_BYTES` -- so how the chunks below that threshold sit among themselves cannot change a task's size. Exactly one property matters: a chunk big enough to be a task on its own must not be picked last and become the straggler the batch waits on. Partitioning those to the front is O(n) in place, and drops the `Vec<u32>` permutation, the `order` field, the size-class table, and the extra load per chunk in `encode_task`. Same measurement as before -- batched `encode`, 64 B documents, corpus repeated to fill 250 ms per pass, 3 interleaved rounds -- shows the two are equivalent in throughput, so the precision the counting sort bought was not worth its code: counting sort partition gpt2/chinese +17.8% +16.9% llama-3/chinese +11.0% +12.9% llama-3/code +6.8% +6.3% llama-3/english +1.5% +5.4% gpt2/code +2.9% +3.0% gpt2/english +4.0% +2.6% Net effect on this PR: 152 insertions become 88, and the worker's inner loop goes back to indexing `chunks` directly.
McPatate
reviewed
Sep 2, 2026
Comment on lines
+478
to
+479
| /// The partition has to leave every chunk that is a task by itself ahead of every chunk that is | ||
| /// not, and it has to keep all of them -- it reorders, it never drops. |
Member
There was a problem hiding this comment.
either reword this or delete pls, it's not very clear
Co-authored-by: Luc Georges <McPatate@users.noreply.github.com>
Co-authored-by: Luc Georges <McPatate@users.noreply.github.com>
ArthurZucker
commented
Sep 2, 2026
ArthurZucker
added a commit
that referenced
this pull request
Sep 2, 2026
#2378 landed the sort -> partition swap that this branch's "take the batch planning off the critical path" commit also made, so `plan_work` conflicted. Both sides had the identical partition loop; the resolution keeps the comment that landed on the base branch and this branch's `seq_start.push(outputs.len())`, which is the other half of that commit and is not on the base. `encode_flat` conflicted only because both sides append to the end of `parallel.rs`; kept whole.
ArthurZucker
added a commit
that referenced
this pull request
Sep 2, 2026
This branch carried three unrelated changes besides the flat batch path. One (the sort -> partition swap) landed as #2378. The other two are now their own PRs and are reverted here so this branch is the flat path alone: * `perf(parallelism): default the pool to physical cores, not SMT siblings` -- 166 lines of pool sizing and Linux `/sys` topology parsing, nothing to do with the batch path. Was #2335. * `perf(pipeline): reuse sequence A's buffer in post_process` -- a change to the general path. Was #2336. Both revert cleanly and the flat path's tests still pass, which is the point of splitting them: each can be reviewed and measured against the base on its own, and a problem with one no longer blocks the other two.
ArthurZucker
added a commit
that referenced
this pull request
Sep 3, 2026
…2378) * perf(pipeline): decide the single-sequence serial fallback before planning `parallel::encode` gated on `inputs.size_bytes() < PARALLEL_MIN_BYTES`, then planned, then bailed to `encode_serial` when planning produced fewer than two tasks. For one document between `PARALLEL_MIN_BYTES` and twice it that bail is guaranteed: `plan_sequence` short-circuits any sequence under `2 * PARALLEL_MIN_BYTES` into a single chunk, one chunk makes one task, and one task fails the check. So the allocations were provably discarded, and the condition is knowable from the length alone. No measurable throughput effect, and it is worth recording why rather than implying one. Because `plan_sequence` short-circuits *before* its added-token scan, the only waste in that window is about six small `Vec` allocations, against ~22 us of encoding for a 9 kB document. Measured with single-document encodes at 4 kB / 9 kB / 20 kB (the first and last as controls, since the guard cannot fire there): every cell within +-1.5%, including the 9 kB window itself. Kept because it removes work that cannot ever be used, not because it is faster. The expensive version of this waste is at 16 kB and above, where `plan_sequence` does run the full `SpecialSegmentIterator` pass over the document, finds no special token, emits one chunk and discards it -- a duplicated added-token scan. This guard does not cover that, and it cannot cheaply: whether a sequence splits is only knowable from that scan. The real fix there is the intra-sequence splitting #2331 left as follow-up, which would make the planning useful instead of merely cheap to skip. * perf(pipeline): order chunks by size class with a counting sort `plan_work` sorted the chunk list with `sort_unstable_by_key(Reverse(len))` and then folded it a second time for the per-sequence `chunk_count`. Two things are wrong with that shape: - `SequenceChunk` is 40 bytes, so a comparison sort does O(n log n) moves of 40 bytes each. The container was never the problem -- `chunks` is a `Vec`, access was already O(1) -- what it sorts is. - It runs in the serial prologue, on the calling thread, before a single worker is spawned. It is pure Amdahl fraction, and worst exactly in the many-small-inputs case tasks were introduced for in #2338: one chunk per sequence means n is the batch size, so an 80k-document batch sorts 80k structs before any encoding starts. The grouping it feeds never needs a total order -- it walks the chunks accumulating bytes and cuts a task every `PARALLEL_MIN_BYTES`. Ordering by power-of-two size class is sufficient, and that is a counting sort: histogram, exclusive prefix sums largest-class-first, scatter. O(n), moving 4-byte indices instead of 40-byte structs, with `chunk_count` coming off the same histogram pass rather than a second fold. `tasks` now index `order` rather than `chunks`, so the worker pays one extra `u32` load per chunk -- against at least `PARALLEL_MIN_BYTES` of tokenizing per task. Batched `encode` of a 5 MB corpus, min of 5 passes with the corpus repeated to fill 250 ms per pass, 3 interleaved rounds, separately built binaries: 64 B documents n_docs base MB/s this commit gpt2/chinese 80457 362.7 +17.8% llama-3/chinese 80457 344.4 +11.0% llama-3/code 32691 709.0 +6.8% gpt2/english 81983 668.4 +4.0% gpt2/code 32691 683.9 +2.9% llama-3/english 81983 663.8 +1.5% 1024 B documents n_docs base MB/s this commit llama-3/chinese 5118 600.9 +2.1% gpt2/chinese 5118 846.5 +1.8% others -1.0% .. +0.9% Positive in every cell where n is large and gone by 1 kB documents, which is the shape an O(n log n) -> O(n) change in a serial prefix should have. The noise floor for this setup is about +-4%, measured from a commit that cannot affect the batch path at all, so the two chinese cells are the solid ones and the sub-3% entries are marginal. Token ids are unaffected: chunks are encoded independently and reassembled by `(seq, idx)`, so visiting order cannot change the result. Verified identical across all four builds in every cell. Only the task boundaries move. * perf(pipeline): partition the chunks rather than counting-sort them Replaces the counting sort from the previous commit with a single partition pass, on review against #2361 which reached the same conclusion independently and expressed it in eight lines. The task grouping never needed the ordering the counting sort provided. Tasks are byte-balanced -- each grown until it holds `PARALLEL_MIN_BYTES` -- so how the chunks below that threshold sit among themselves cannot change a task's size. Exactly one property matters: a chunk big enough to be a task on its own must not be picked last and become the straggler the batch waits on. Partitioning those to the front is O(n) in place, and drops the `Vec<u32>` permutation, the `order` field, the size-class table, and the extra load per chunk in `encode_task`. Same measurement as before -- batched `encode`, 64 B documents, corpus repeated to fill 250 ms per pass, 3 interleaved rounds -- shows the two are equivalent in throughput, so the precision the counting sort bought was not worth its code: counting sort partition gpt2/chinese +17.8% +16.9% llama-3/chinese +11.0% +12.9% llama-3/code +6.8% +6.3% llama-3/english +1.5% +5.4% gpt2/code +2.9% +3.0% gpt2/english +4.0% +2.6% Net effect on this PR: 152 insertions become 88, and the worker's inner loop goes back to indexing `chunks` directly. * Apply suggestion from @McPatate Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> * Apply suggestion from @McPatate Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> * Apply suggestion from @ArthurZucker --------- Co-authored-by: Luc Georges <McPatate@users.noreply.github.com> (cherry picked from commit cbd94b8)
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.
plan_worksorted the chunk table by descending length in the serial prologue —O(n log n) moves of a 40-byte struct on the calling thread, where n is the batch size
when each sequence is one chunk. A single partition pass is O(n) and keeps the only
property the scheduler needs. +2.6% to +16.9% on 80k-document batches, ids identical.
Relationship to #2361
#2361 already contains this, reached independently, and its eight-line version is
better than the counting sort I first wrote — so this PR now uses that shape. Whichever
lands first, the other should drop the hunk.
The reason to keep this one separate: #2361 bundles three changes, and #2365 undercuts
two of them.
encode_batch_flat(+222 inmod.rs, +167 inparallelism.rs) and theconcat-to-pool machinery exist so workers share nothing — but #2365 shows the shared
ScratchPoolmutex was the contention, taking scaling efficiency from 3.3–12.3% to69–82%. #2361's "3.4× over the general path" was measured against a baseline whose
contention #2365 removes, so that comparison needs redoing before the flat path earns
its keep.
This PR is the planning fix alone: one file, +88/−2.
The change
Tasks are already byte-balanced — each grown until it holds
PARALLEL_MIN_BYTES— so howthe sub-threshold chunks sit among themselves cannot change a task's size. Only one thing
matters: a chunk big enough to be a task on its own must not be picked last and become the
straggler the batch waits on.
Measured
Batched
encode, 64 B documents, corpus repeated to fill 250 ms per pass (borrowed frommeasure_scaling, for its reason — one call is a few ms and spawns up tocurrent_num_threads()rayon tasks), min over passes, 3 interleaved rounds, separatelybuilt binaries.
Gone by 1 kB documents, which is the shape an O(n log n) → O(n) change in a serial prefix
should have.
ids_hashidentical to base in every cell.The middle commit is the counting-sort version and the last one removes it; the history is
kept so the comparison is on record, and it squashes to the eight-line form.
The other commit
The single-sequence early-out has no measurable effect and its message says so:
plan_sequenceshort-circuits before its added-token scan, so the only waste in the8–16 KiB window is ~6 small allocations. Measured at 4 kB / 9 kB / 20 kB single documents
(first and last as controls, where the guard cannot fire): every cell within ±1.5%. Kept
only because it removes provably-discarded work. Drop it if you would rather have the
smaller diff.
Note on measuring this
Engine::encodein tokbench takes a single&str, soplan_workthere only ever sees onechunk, fails
tasks.len() < 2, and falls back toencode_serial— the planner isunreachable through that trait, and
--scalingruns N independent single-document engines.These numbers come from a batched harness. A batch method on
Enginewith a default impllooping the single-document one would make this shape measurable in tokbench proper.