perf(r3): accelerate routed-expert transport with packed arrays

[dyurk-lila]

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

• fix(r3): align routed-expert metadata and padding #1908 — fix(r3) align routed-expert metadata and padding

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

Problem

The routed-expert (R3) generation path returns one expert ID for every generated token, MoE layer, and top-k slot. Serializing that integer tensor as nested JSON lists creates a large HTTP payload and substantial Python/JSON overhead before training sees the data.

Implementation

• Compact routed-expert IDs to the smallest safe uint8, int16, or int32 NumPy dtype.
• Base64-encode the contiguous buffer alongside its shape and dtype, and decode it back into a validated compact NumPy array on the client.
• Serialize the vLLM response and parse client responses with orjson.
• Gate packed routes behind the existing routed-expert request flag, so the transport is only used when routed experts are requested.
• Accept the decoded NumPy route arrays at preprocessing, validating dtype and shape before padding and tensor conversion.

This change is intentionally scoped to transport and preprocessing: it does not alter batching, Experience, device placement, or backend replay setup.

Testing

CPU tests run with:

uv run --isolated --extra dev --extra fsdp pytest \
  tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py \
  tests/backends/skyrl_train/inference_servers/test_routed_experts_wire.py \
  tests/train/dataset/test_preprocess.py \
  tests/train/generators/test_generator_output_utils.py \
  tests/train/generators/test_skyrl_gym_generator.py \
  tests/train/test_trainer_utils.py -q --ignore=tests/backends/skyrl_train/gpu

The suite covers route preprocessing, dtype selection, base64 round trips, malformed payloads, and the packed wire codec. ruff and black (line length 120) pass on the changed files.

This is part of a small series of routed-expert transport improvements.

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.