[faresobeid]

Reworked onto vLLM 0.31 (stacked on #3888)

Branch feat/payload-segments now holds one commit 1bcca4115 on top of #3888's head e89019444 (squashed from the six earlier commits; the content is the same as 99fa8c814, with only an import conflict in batch.py resolved). This PR is closed, so GitHub still shows the old head and base main. The intended base is chore/bump-vllm-0.31.

• Checks on 0.31 (/home/fares/scratch/venv031, vLLM 0.31.0): pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/inference: 307 passed on the stack top (feat: by-handle payload fields table (TRR weights, score-centering logprobs) #3836). Ruff: clean.
• GPU (2026-10-06, vLLM 0.31.0, Qwen3-30B-A3B, BI): rows written by a feat: write router-replay ids and sampling masks by handle on multi-node runs #3820 server, pushed through prepare_batch + DataLoader._micro_batch_to_tensor (seq_len 16384/4096 with truncation, 2 DP, CP 1/2/4): 140/140 micro-batch comparisons bitwise equal to inline. The earlier None-vs-all(-1) mask difference no longer shows up.

---

What

This is the consumer half of passing per-token side arrays (router-replay ids, sampling masks) by handle (design: payload-by-handle-design.md). It adds no producer and changes no behavior on its own. #3820 (stacked) adds the vLLM-side writer, the switch and the GC.

• prime_rl/transports/payload.py:
  • PayloadSegment(field, file, offset, pos, rows, dtype, shape) (msgspec, array-like): rows fixed-width rows of one field in a shared file, starting at token position pos.
  • clip_segments(segments, lo, hi).
  • read_field(segments, field, lo, hi, fill): preads only the rows in [lo, hi) into an int32 array filled with fill. Rows narrower than the field's widest row fill the leading entries (sampling masks are written at each request's own width). Later segments win.
• Transport: TrainingSample.payload / MicroBatch.payload are appended last (wire-stable). They coexist with the inline routed_experts / sampling_mask.
• Packing (orchestrator/batch.py): truncation clips the segments, packing shifts them by the sample's offset in the micro batch, padding is "no segment", router-replay presence still gates bin compatibility (inline or by handle), dummy batches drop mask segments.
• Trainer (trainer/rl/data.py, train.py):
  • A text micro batch with segments is read by each rank for its own CP chunk only. Routing reads [lo, hi). The mask reads [lo+1, hi+1), so it is already on label positions (shift_tensor_left). train.py skips the shift and shard for those tensors (payload_cp_window).
  • Multimodal micro batches read the full sequence, because the model may defer CP.
  • Fills: experts 0, mask -1. The mask width is trimmed to the window's widest row (inline uses the micro batch's widest row).

Validation

• test_payload_by_handle_matches_inline (one new test). 24 random samples carry router-replay ids and ragged sampling masks inline, and the same rows by handle: routing split across 2 files, masks one segment per position at its own width. Both go through prepare_batch (truncation at 32, 2 DP ranks, padding to the CP degree, dummies) and DataLoader._micro_batch_to_tensor. The tensors are bitwise identical to inline for CP=1. For CP=2 they equal shard_for_cp(shift_tensor_left(inline)) on every rank. The by-handle mask is narrower per window, and every dropped column is -1.
• pytest tests/unit/test_configs.py tests/unit/orchestrator tests/unit/train/rl -m "not gpu": 297 passed. ruff check and ruff format --check pass.
• Synthetic huge batch on CPU (64 × 32k-token samples, Qwen3-30B-A3B routing 48×8 uint8, 30 turns each, masks ≤64 ids on 20% of tokens, 8 DP × CP 4), inline vs by handle:

	inline	by handle
orchestrator: samples resident	907 MiB	81 MiB
orchestrator: peak during pack + encode	2,895 MiB	332 MiB
ZMQ payload, all DP ranks	895 MiB	68 MiB (all of it token lists/floats; segments ≈ 0.4 MiB)
egress (× 4 CP subscribers)	3,581 MiB	273 MiB
pack + encode time	3.2 s	2.3 s
DP rank 0 decode + materialize (4 CP ranks)	1.07 s	0.86 s (each rank reads ¼ of the rows)

Not done here

• Producer, verifiers/renderers plumbing, the switch and the payload-root GC: feat: write router-replay ids and sampling masks by handle on multi-node runs #3820.

GPU validation (2026-10-04)

Setup: 8x H200, vLLM 0.30.0, one integration branch (#3832 ⊃ #3831 ⊃ #3830, + #3794, #3793, #3820 ⊃ #3819) with router#61 + router#60. Exact comparisons use VLLM_BATCH_INVARIANT=1 (BI) with Qwen/Qwen3-30B-A3B (TP1, 48 layers, top-8) and the same token ids and seeds (temperature 1, top_p 0.95, top_k 50).

• By-handle vs inline trainer tensors (colocated BI server, 16 requests of 500-8000 prompt tokens, same seeds with and without payload_dir): tokens identical; file rows equal inline routing (uint8 [rows, 48, 8]) and inline masks for 16/16. Real rows through prepare_batch + DataLoader._micro_batch_to_tensor (seq_len 16384 and 4096 with truncation, 2 DP ranks, CP=1/2/4): routing equals shard_for_cp(inline) and masks equal shard_for_cp(shift_left(inline)) in 140 micro-batch/rank comparisons.
• One difference, fixed here: in 49 windows (samples truncated before any completion token) by-handle gave None while inline gave an all -1 mask. Commits b1cee0f + 99fa8c8: when truncation cuts all mask rows, prepare_sample keeps a zero-row, width-1 sampling_mask placeholder, so the trainer gets the same all -1 width-1 tensor as inline. test_payload_by_handle_matches_inline now includes such a sample (fails without the fix).
• Not run: a CP=4 trainer run (CP=4 covered only through the DataLoader windows on CPU), crash/resume GC, payload FS throughput at GLM-5.3 scale.

Landing order

1. feat(client): pass through by-handle payload segments renderers#171 and feat(v1): by-handle payload segments on traces verifiers#2765
2. feat: by-handle payload segments: reader, packing, CP-window reads, GC #3819 (reader/packing, no behavior change)
3. feat: write router-replay ids and sampling masks by handle on multi-node runs #3820 (writer, switch and GC; re-pin deps/renderers / deps/verifiers to the merged commits first)
4. feat(pd): stitch by-handle routed_experts payload segments router#60 (independent), then let rl turn by-handle on for disaggregated P/D

---

Note

Medium Risk
Changes how MoE router replay and sampling masks reach the trainer under context parallelism; incorrect windowing or sharding would silently skew training, though inline paths are unchanged and a parity test covers the new reader.

Overview
Adds by-handle transport for per-token side data (routed_experts, sampling_mask): inference will write rows to shared files and pass lightweight PayloadSegment handles on TrainingSample / MicroBatch, alongside the existing inline fields.

New transports/payload.py defines segments plus clip_segments and read_field (pread into token windows). The orchestrator clips and repositions segments on truncation, shifts them when packing micro-batches, treats payload router replay like inline for bin compatibility, and strips mask segments from dummy padding batches. The trainer DataLoader takes CP rank/size; when payload is present it materializes experts/masks from files—for text batches only each rank’s CP chunk, with masks already on label positions—and sets payload_cp_window so train.py skips the usual mask shift and CP shard for those tensors. Multimodal batches still read the full sequence when CP may be deferred in the model.

A unit test asserts by-handle tensors match the inline path through pack/truncate/pad and CP=1/2.

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