Skip to content

perf(pipeline): partition plan_work's chunks instead of sorting them - #2378

Merged
ArthurZucker merged 6 commits into
feat/train_encode_splitfrom
perf/plan-work-bucketing
Sep 2, 2026
Merged

ArthurZucker merged 6 commits into
feat/train_encode_splitfrom
perf/plan-work-bucketing

Conversation

@ArthurZucker

@ArthurZucker ArthurZucker commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

plan_work sorted 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 in mod.rs, +167 in parallelism.rs) and the
concat-to-pool machinery exist so workers share nothing — but #2365 shows the shared
ScratchPool mutex was the contention, taking scaling efficiency from 3.3–12.3% to
69–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

let mut boundary = 0;
for i in 0..chunks.len() {
    if chunks[i].range.len() >= PARALLEL_MIN_BYTES {
        chunks.swap(boundary, i);
        boundary += 1;
    }
}

Tasks are already byte-balanced — each grown until it holds PARALLEL_MIN_BYTES — so how
the 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 from
measure_scaling, for its reason — one call is a few ms and spawns up to
current_num_threads() rayon tasks), min over passes, 3 interleaved rounds, separately
built binaries.

cell n_docs base MB/s partition
gpt2/chinese 80457 361.5 +16.9%
llama-3/chinese 80457 337.7 +12.9%
llama-3/code 32691 701.3 +6.3%
llama-3/english 81983 651.4 +5.4%
gpt2/code 32691 675.9 +3.0%
gpt2/english 81983 668.3 +2.6%

Gone by 1 kB documents, which is the shape an O(n log n) → O(n) change in a serial prefix
should have. ids_hash identical 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_sequence short-circuits before its added-token scan, so the only waste in the
8–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::encode in tokbench takes a single &str, so plan_work there only ever sees one
chunk, fails tasks.len() < 2, and falls back to encode_serial — the planner is
unreachable through that trait, and --scaling runs N independent single-document engines.
These numbers come from a batched harness. A batch method on Engine with a default impl
looping the single-document one would make this shape measurable in tokbench proper.

…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.
@ArthurZucker ArthurZucker changed the title perf(pipeline): order plan_work's chunks with a counting sort perf(pipeline): partition plan_work's chunks instead of sorting them Sep 2, 2026
Comment thread tokenizers/tk-encode/src/tokenizer/pipeline/parallel.rs Outdated
Comment thread tokenizers/tk-encode/src/tokenizer/pipeline/parallel.rs Outdated
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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

either reword this or delete pls, it's not very clear

ArthurZucker and others added 2 commits September 2, 2026 11:24
Co-authored-by: Luc Georges <McPatate@users.noreply.github.com>
Co-authored-by: Luc Georges <McPatate@users.noreply.github.com>
Comment thread tokenizers/tk-encode/src/tokenizer/pipeline/parallel.rs Outdated
@ArthurZucker
ArthurZucker merged commit cbd94b8 into feat/train_encode_split Sep 2, 2026
12 of 15 checks passed
@ArthurZucker
ArthurZucker deleted the perf/plan-work-bucketing branch September 2, 2026 09:28
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)
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