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

[dyurk-lila]

Note on the diff: This is the bottom of the routed-expert-replay / sampler-support series; it applies directly to main. Later PRs in the series stack on top of it.

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.

Also included: non-finite sampled logprobs

The /skyrl/v1/generate endpoint previously raised on any non-finite logprob, on the
assumption that orjson would refuse to serialize it. orjson actually emits JSON
null, so the guard was unnecessary — and vLLM does report a non-finite logprob for a
token it just sampled (roughly one per 42 minutes of rollouts at 192-GPU scale).
FastAPI maps the ValueError to HTTP 400, which callers treat as a fatal invariant
violation, so a single token could tear down an entire run.

Non-finite values are now floored to -9999.0, matching what vLLM itself does at its
serving boundaries, with the pre-existing missing-entry branch folded into the same
path. Note isfinite also catches NaN, which vLLM's own max() floor lets through
because max returns its first argument on a False comparison. The loop is extracted
into a vLLM-free helper (logprobs_wire.py) so it is unit-testable without a live
engine, and logs a warning identifying the request and clamp count.

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/backends/skyrl_train/inference_servers/test_logprobs_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, the packed wire codec, and clamping of missing/-inf/NaN sampled logprobs. 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.

[erictang000]

/gemini review

[gemini-code-assist]

Code Review

This pull request optimizes the serialization and transmission of routed expert indices and logprobs between the vLLM server and the remote inference client by introducing a binary payload contract using base64-encoded NumPy arrays. Key feedback includes safely detaching and moving CUDA tensors to the CPU before NumPy conversion in the server actor, using more robust dictionary retrieval for logprobs to prevent potential attribute errors, and relaxing shape dimension type checks to accommodate NumPy integer types.