perf(r3): collate routes through a container-aware pool, and refuse replay under VPP

[erictang000]

Note

Standalone-reviewable reopen of #2078 — original work by @dyurk-lila.
Commits are pushed unchanged, so author and committer metadata are preserved.
The branch lives on NovaSky-AI/SkyRL rather than a fork so CI runs.

Stack position 3/11 — based on #2232 (branch dyurk/r3-incremental-traces).
Because the base is the previous branch in the stack, the Files changed tab above shows only this PR's own changes, not the cumulative stack. Merge bottom-up.

Supersedes #2078.

---

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub cannot show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor assemble routed-expert traces incrementally

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

Problem

Routed-expert replay arrives as one variable-length NumPy array per trajectory. Collating those arrays into the left-padded global training batch is dominated by allocating and first-touching a large route buffer. The previous serial path also rebuilt padding intermediates and rescanned route IDs even though the inference boundary had already selected a canonical integer dtype.

Separately, replay is unsafe with Megatron virtual pipeline parallelism. Each forward appends a microbatch to every local replay FIFO, while an interleaved model chunk consumes only its own entry during backward. The queues can therefore desynchronize without an immediate crash.

Implementation

• Add affinity- and cgroup-aware CPU discovery for both cgroup v2 and v1.
• Size controller-side pools from the usable CPU allocation, reserving capacity for colocated services.
• Fill disjoint trajectory regions through a bounded thread pool; worker exceptions are propagated to the caller.
• Allocate the final route tensor once, write canonical route arrays directly, and fill only left-padding and uncaptured suffix rows with valid dummy routes.
• Reuse the CPU-topology helper for checkpoint delta workers, where the operation is already a barrier and does not need a core reserve.
• Reject routed-expert replay with virtual pipeline sizes greater than one after Megatron configuration resolution, so provider defaults cannot bypass validation.

The sender remains responsible for route compaction. Collation validates the accepted uint8, int16, and int32 dtypes and promotes a mixed batch to the widest input dtype without another value scan.

Performance

The included packed-route collation benchmark compares identical outputs across serial and pooled implementations. On a 32-core host with OMP_NUM_THREADS=1, a representative 40-layer global batch improved from 29.61 s to 3.35 s (median of five runs).

The result depends on the controller running with one OpenMP thread, as it does under its one-CPU Ray allocation. With unrestricted library threads, the pool can oversubscribe the container and lose performance; the worker cap and cgroup-aware sizing make that dependency explicit.

Testing

• CPU-topology tests cover affinity, cgroup v1/v2 quotas, unlimited and malformed limits, core reservation, and pool caps.
• Parallel-fill tests cover serial and parallel execution, disjoint ownership, worker limits, empty inputs, and exception propagation.
• Preprocessing tests cover route padding, dtype preservation and promotion, non-contiguous arrays, mismatched shapes, and invalid route lengths.
• Configuration tests cover both explicit and provider-resolved virtual pipeline settings.

---

Note

Medium Risk
Changes MoE rollout→train route bookkeeping and batch collation (correctness-sensitive) plus Megatron replay config guards; parallel CPU pools affect performance but not model math.

Overview
Routed-expert (R3) replay now stitches multi-turn rollouts incrementally: RoutedExpertTrace / TokenMetadataTrace accumulate per-token route rows, the gym generator sends routed_experts_prompt_starts on each LLM call, and RemoteGenerateClient owns the HTTP generate path (including routed_experts_prompt_start in sampling params). Collation no longer rescans or re-compacts routes at batch time—it validates canonical uint8/int16/int32 arrays and packs them in parallel.

Performance: New cpu_topology sizes thread pools from process affinity and cgroup CPU limits; fill_batch_rows fills disjoint trajectory slices in a bounded pool (with core reserve for Ray). Route tensors are allocated once and written directly, with torch replay_padding_row for left/right padding; the NumPy padding helper is removed. Delta checkpoint apply/publish pools use the same helper with zero reserve.

Safety: MoE routing replay is rejected when virtual pipeline parallelism is > 1 (config validation and Megatron provider resolution), because interleaved chunks can desync RouterReplay’s backward FIFO.

^{Reviewed by Cursor Bugbot for commit 8d95e20. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces a new CPU topology utility to dynamically determine usable CPU cores based on process affinity and cgroup quotas, which is then used to size thread pools for weight synchronization and parallel batch filling. It also refactors routed expert index collation to perform parallel batch filling and adds validation to prevent the incompatible combination of Megatron routing replay and virtual pipeline parallelism. A review comment suggests catching OSError in addition to AttributeError when calling os.sched_getaffinity to prevent crashes in restricted container environments.

[greptile-apps]

Confidence Score: 4/5

The PR should not merge until the configuration guard permits virtual_pipeline_model_parallel_size=1, matching the resolved-provider check and the stated interleaving constraint.

Findings

1.  VPP Size One Rejected ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Canonical NumPy route arrays] --> B[Validate dtype, shape, and route length]
  B --> C[Select widest batch dtype]
  C --> D[Allocate final Torch route tensor once]
  D --> E[Container-aware bounded thread pool]
  E --> F1[Fill trajectory 1 and padding]
  E --> F2[Fill trajectory 2 and padding]
  E --> FN[Fill trajectory N and padding]
  F1 --> G[Global training batch]
  F2 --> G
  FN --> G

Loading

_{Reviews (1) · Last reviewed commit: "style: tighten route collation comments ..."}

[erictang000]

Benchmarked the collation path: the win is the single-write change, not the pool

(Updated: the original version of this review reported speedup ratios, which oversold the thread pool. Absolute per-step times below, and the pool has now been removed.)

This runs once per training step (already timed as convert_to_training_input), so absolute cost per call is the figure that matters. Benchmarked on a 192-core host; fresh buffer per rep so first-touch faults are re-incurred as in production; median of 5-7 reps; torch.equal(old, new) asserted before timing.

Absolute cost of one collation (ms)

buffer	pre-PR	after single-write change (serial)	+ pooled(8)	+ pooled(32, original cap)
0.40 GB	398.7	70.8	48.0	107.6
1.61 GB	1596.4	206.5	86.6	147.4
4.09 GB	2568.3	244.3	150.9	204.8
16.1 GB	-	721.3	528.6	607.8
32.2 GB	-	1196.0	1126.0	1135.0

Saving per step, by change

change	0.40 GB	1.61 GB	4.09 GB	16.1 GB	32.2 GB
drop per-sample compact_routed_expert_indices	18 ms	76 ms	467 ms	-	-
replace full-buffer prefill with torch.empty + exact-region writes	310 ms	1314 ms	1857 ms	-	-
thread pool, best case	25 ms	124 ms	92 ms	193 ms	70 ms

Replacing make_replay_padding_indices_np's prefill-then-overwrite with single-write region fills saves 0.3-1.9 seconds per step. That is the PR, and it is worth landing on its own.

The pool saves 25-200 ms, and the margin shrinks as buffers grow -- at 32 GB it is 70 ms out of 1196 ms, because by then the fill is memory-bandwidth bound rather than page-fault bound. So the large-buffer case that motivates the pool is where it helps least.

The pool was also mis-sized

workers	0.40 GB	1.61 GB	4.09 GB	16.1 GB	32.2 GB
1	1.00x	1.00x	1.00x	1.00x	1.00x
4	1.38x	1.94x	1.60x	1.38x	1.20x
8	1.52x	2.43x	1.61x	1.36x	1.06x
16	1.06x	1.86x	1.42x	1.30x	1.06x
32 (original cap)	0.68x	1.43x	1.19x	1.19x	1.05x
64	0.27x	0.62x	0.60x	-	-

At the original cap of 32 the pool was slower than a serial fill below ~1 GB (107.6 ms vs 70.8 ms) -- thread startup exceeded the fill.

What I pushed

• perf(r3): drop the batch-fill thread pool -- removes parallel_fill.py and its tests; the fill is now a plain loop. A tenth of a second per step does not pay for two modules, a per-step ThreadPoolExecutor, and ~130 lines of tests.
• skyrl/utils/cpu_topology.py is kept. It has a second consumer: the weight-sync apply and publish pools in weight_sync/delta/checkpoint.py, where pool_workers replaced min(32, os.cpu_count()) and correctly respects the cgroup quota. On this host that is a real difference -- cpu.max = "9140000 100000" gives a 91-CPU quota against 192 visible cores.
• test: allow virtual_pipeline_model_parallel_size=1 under routing replay -- this test has been failing since f9bdc69. That commit relaxed the assert to vpp_size is None or vpp_size <= 1 (matching megatron_worker.py, which only raises above 1), but test_routing_replay_refuses_virtual_pipeline_parallelism[1] still asserted 1 is refused. Moved 1 to the allows-parametrize.

One unrelated failure remains, test_megatron_validation_rejects_mxfp8_with_fp8_param, which also fails on main -- not from this PR.

Note for whoever propagates this up the stack

fill_batch_rows has call sites in later PRs (#2235 adds a second, #2238 a third). Those need the same one-line conversion to a plain loop when this merges up. I stopped at this branch rather than rebasing the descendants, since the stack is being propagated with merge commits.

Remaining, non-blocking

1. The read-only fallback sample_indices.copy(order="C") is a full extra copy per sample. It should not fire today (generator output is built in-process and writable), but it is silent if provenance changes. Writing through a numpy view of the destination avoids it -- numpy does not require a writable source, only torch.from_numpy does.
2. The int32 warning fires per batch, i.e. every step.
3. The VPP guard is unrelated to route collation; fine to keep given its size, but it does make the PR two things.