[mikasenghaas]

Summary

Logs the trainer's batch load at info and trims the top-logprob part of it. Output tensors are unchanged.

• Timing at info, per node: time waiting for the batch, decode vs. convert time with micro-batch and token counts, and time from load to the first micro batch on the GPU. The last one includes the global loss-scale all-reduce, so a rank still loading its batch shows up on every other rank.
• Zero-copy top-logprob and sampling-mask decode. TopLogprobs and SamplingMask buffers decode as memoryviews into the received frame instead of bytes copies. Both are copied into their [tokens, width] grids during conversion, so no view outlives get_batch. RoutedExperts.data stays bytes: routing is kept as a tensor for the whole step, and a view would keep the entire frame alive with it.
• Top-logprob grid. With a single head (no masked tokens), the ids and logprobs are read in place instead of going through the two-head interleave. When every row carries either the full head or nothing (sampled vs. context tokens, the usual case), they are copied row-wise into the grid instead of through an element mask. Ragged rows keep the existing scatter.

Measurements

One 262,144-token micro batch, decoded from its msgpack frame and converted single-threaded (OMP_NUM_THREADS=1, as under torchrun). Median of 15, three repeats, main @ 14ef914 vs this branch on the same machine:

	main	this PR
router replay (43 x 6 uint8) + top-128 on 40% of tokens, 181 MB frame	0.49-0.50 s / 1M tokens	0.45-0.50 s / 1M tokens
score-centering sampling masks (128 ids + logprobs) on 40% of tokens, 114 MB frame	0.67-0.69 s / 1M tokens	0.65-0.69 s / 1M tokens

Most of the earlier load-time gain landed on main with #4226. What is left here is about 5-8% on router-replay batches, none measurable on score-centering batches, and the timing logs.

Verification

• _micro_batch_to_tensor output compared tensor by tensor against main on four frames (router replay with full-width and with ragged top-k rows; score-centering masks with full-width and with ragged rows): 10 tensors each, all equal in value, dtype and shape.
• tests/unit/transports, tests/unit/orchestrator and tests/unit/train/rl (CPU) pass.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches RL batch decoding and top-logprob tensor assembly on the training hot path; behavior is intended to be identical but mistakes could skew replay or masks.

Overview
Speeds up trainer batch decode/convert (~2.7× on router-replay + top-k workloads) without changing output tensors, and surfaces load timing at info level.

Wire structs SamplingMask and TopLogprobs now decode buffer fields as memoryviews into the received ZMQ frame (zero-copy); RoutedExperts.data stays bytes so routing tensors do not pin the whole frame. get_batch logs micro-batch count, token count, and separate decode vs. convert durations; train.py promotes batch-wait logging to info and logs time until the first micro batch hits the GPU (including cross-rank loss-scale sync).

Top logprobs assembly adds a single-head fast path and, when every row is either full-width or empty, row-wise reshape into the [tokens, width] grid instead of masked scatter; ragged rows keep the prior scatter behavior.

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

[mikasenghaas]

Closing: most of this PR's gain already landed on main with #4226, and what is left isn't worth the change.

Re-measured against current main (14ef914), one 262,144-token micro batch decoded and converted single-threaded, median of 15 over three repeats:

	main	this PR
router replay + top-128	0.49-0.50 s / 1M tokens	0.45-0.50 s / 1M tokens
score-centering sampling masks	0.67-0.69 s / 1M tokens	0.65-0.69 s / 1M tokens

• The 2.7x in the description was measured against 1499de6, before perf: per-token streams as numpy arrays, micro batches as raw buffers #4226. perf: per-token streams as numpy arrays, micro batches as raw buffers #4226 made the per-token streams raw buffers, decoded as memoryviews and converted with one numpy copy.
• main already keeps uint8 routing on the host and widens it on the GPU.
• Zero-copy routing had to be reverted: the routing tensor lives for the whole step, so a view into the frame would keep the whole frame alive.
• That leaves about 5-8% on router-replay batches, nothing measurable on score-centering batches, and the info-level load timing. If we want the timing lines, they can go in as a small separate PR.