perf(sample-support): add packed base64 transport by dyurk-lila · Pull Request #1915 · NovaSky-AI/SkyRL

dyurk-lila · 2026-07-16T22:30:07Z

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:

perf(r3): accelerate routed-expert transport with packed arrays #1909 — perf(r3) packed-array transport

perf(r3): cast expert indices to int32 for pp local layers only #1910 — perf(r3) expand routes to local layers only

refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator

refactor(generators): assemble routed-expert traces incrementally #1912 — refactor incremental routed-expert traces

feat(megatron): add bounded sampler-support replay #1913 — feat(megatron) bounded sampler-support replay

feat(fsdp): add bounded sampler-support replay #1914 — feat(fsdp) sampler-support replay

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

Problem

Dense sampler-support rows are large integer arrays. Returning them as nested JSON lists materializes one Python integer per candidate and spends substantial controller time formatting decimal JSON.

Implementation

keep vLLM's extracted [tokens, top_k] support as a contiguous int32 NumPy array;
encode that buffer with pybase64 inside the existing orjson response envelope;
make PackedSampleSupportSet the only wire response contract, with strict dtype, shape, byte-count, ID-range, and trailing-padding validation;
decode exactly once inside RemoteInferenceGenerator before populating RemoteGenerateResult.sample_support;
let both external generator callers and RemoteInferenceClient consume the same logical decoded result.

This change affects only the physical rollout transport. The trainer-side representation and replay semantics are unchanged.

Performance

The included benchmark (skyrl/benchmarks/bench_dense_sample_support_wire.py) uses orjson on both paths and compares equivalent decoded outputs. For 4096-token responses:

Support	Packed/list bytes	Encode speedup	Decode speedup
top-k 8, full	0.82x	10.8x	1.03x
top-k 8, variable top-p	1.05x	10.0x	0.97x
top-k 64, full	0.85x	20.6x	1.22x
top-k 64, variable top-p	1.14x	15.4x	1.11x

The packed envelope is primarily a controller-CPU optimization. Sparse top-p rows can be slightly larger because base64 still carries -1 padding.

Testing

new codec tests cover round-trip, malformed payloads, empty token dimensions, legacy-list rejection, and shape/size/dtype validation;
updated client and vLLM endpoint tests verify packed decode through RemoteGenerateResult and the real /skyrl/v1/generate response;
uv run --isolated --extra dev --extra fsdp pytest over the changed inference-server test files passes (76 tests);
Ruff and Black pass on the changed files.

gemini-code-assist · 2026-07-16T22:33:30Z

Warning

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

Serializing routed-expert IDs as nested JSON lists creates a large HTTP payload and substantial Python/JSON overhead before training sees the data. - Compact routed-expert IDs to the smallest safe uint8, int16, or int32 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. - Accept the decoded NumPy route arrays at preprocessing, validating dtype and shape before padding and tensor conversion. Because orjson emits JSON null rather than raising for non-finite floats, the sampled-token logprob path floors missing and non-finite values to -9999.0 -- the same floor vLLM applies at its own serving boundaries. This keeps one bad logprob from failing the whole generate request. Note the isfinite test also catches NaN, which vLLM's max(logprob, -9999.0) misses since max returns its first argument on a False comparison. This change is intentionally scoped to transport and preprocessing: it does not alter batching, Experience, device placement, or backend replay setup.

Under pipeline parallelism each rank replays only its local router layers, but the Megatron worker eagerly expanded the full global-layer routed-expert tensor to int32 before replay setup, allocating a large device temporary for unused layers. Keep routed-expert IDs in their compact dtype through whole-batch device movement, index_select the current PP stage's router layers before metadata alignment, and perform the single int32 conversion inside _split_replay_indices so only the bounded PP-local slice is materialized as int32. Also validate the 4D replay-indices shape up front. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteInferenceGenerator and a RemoteGenerateResult dataclass. RemoteInferenceClient now owns an internal generator and delegates session management, _post, and _generate_single to it. This is a pure refactor with no functionality change: endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are all preserved. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Add the core of bounded sampler-support replay across data plumbing, inference capture, and the Megatron dense-path trainer. - Inference: /skyrl/v1/generate gains a return_sample_support flag that requests flat_logprobs from vLLM; the fixed-width rows are decoded into per-token support IDs and surfaced through RemoteInferenceGenerator, RemoteGenerateResult, and InferenceEngineOutput.rollout_sample_support. - Data plumbing: support rows ride the shared TokenMetadataTrace through multi-turn accumulation, dynamic filtering, truncation, replay buffering, and microbatch padding; preprocessing validates support width, trailing -1 padding, int32 vocab IDs, and presence of every loss-bearing token. - Config: reject sampler settings whose rollout distribution cannot be replayed exactly (temperature > 0, top_k > 1, repetition penalty 1.0, no arbitrary additional_kwargs). - Megatron: renormalize sampled-token scores over recorded support with fixed-shape gathers; fused LM-head path projects only selected candidate pairs and chunks by logprobs_chunk_size; controller-packed rows reuse TokenMetadataLayout. New utils/sample_support_replay.py holds the core. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…om its top-k support (NovaSky-AI#66) * fix(sample-support): repair rows where the sampled token is absent from its top-k support vLLM's approximate top-k/top-p pivot can let a sampled survivor rank just beyond top_k, so it is missing from the captured support row and the downstream sample-support invariant hard-crashes the run. When the sampled id is absent, overwrite the weakest support member with it, preserving width==top_k, single-occurrence, and trailing padding. Emit one aggregated warning per call. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * style(sample-support): apply black formatting to regression test --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Extend bounded sampler-support replay to the FSDP backend so policy and reference forwards apply the same support-conditioned normalization as Megatron, keeping recomputed logprobs on the same distribution. - Enable replay for trainer.strategy=fsdp with the existing sampler-support constraints; thread support IDs and the response loss mask through policy training, policy recomputation, and reference forwards only when enabled. - Right-align the response loss mask into the full token layout and keep sampled IDs, support IDs, and the mask aligned through padding removal, next-token shifting, and Ulysses SP slicing; reuse padding-removal indices. - Generalize Ulysses input partitioning to [batch, seq, ...] metadata with a configurable padding value, avoiding a second validity-mask partition. - Share the backend-neutral aligned replay helper for support normalization and synthetic-EOS handling; preserve feature-disabled behavior. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Transport dense sampler-support vocab IDs as a base64-encoded int32 NumPy buffer inside the orjson response envelope, replacing nested JSON integer lists. - keep vLLM's [tokens, top_k] support as a contiguous int32 array and encode it via pybase64 in the response envelope; - add sample_support_set_wire with PackedSampleSupportSet as the wire contract, validating dtype, shape, byte count, ID range, and trailing -1 padding; - decode once in RemoteInferenceGenerator before populating RemoteGenerateResult.sample_support; - add a benchmark and tests covering the codec, malformed payloads, empty token dimensions, legacy-list rejection, and endpoint decode. Trainer representation and replay semantics are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

dyurk-lila force-pushed the upstream/sample-support-packed-wire branch 2 times, most recently from c583c79 to 1d98672 Compare July 20, 2026 20:25

erictang000 self-assigned this Jul 20, 2026

dyurk-lila force-pushed the upstream/sample-support-packed-wire branch 2 times, most recently from ea66ff0 to c358853 Compare July 29, 2026 16:09

dyurk-lila force-pushed the upstream/sample-support-packed-wire branch from c358853 to 1eabb34 Compare July 29, 2026 21:07

dyurk-lila and others added 8 commits July 31, 2026 22:49

test(generators): drop obsolete turn route argument

98a6724

dyurk-lila force-pushed the upstream/sample-support-packed-wire branch from 1eabb34 to 187044b Compare July 31, 2026 22:52

Navigation Menu

Uh oh!

perf(sample-support): add packed base64 transport - #1915

perf(sample-support): add packed base64 transport#1915
dyurk-lila wants to merge 9 commits into
NovaSky-AI:mainfrom
dyurk-lila:upstream/sample-support-packed-wire

dyurk-lila commented Jul 16, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dyurk-lila commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Implementation

Performance

Testing

Uh oh!

gemini-code-assist Bot commented Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dyurk-lila commented Jul 16, 2026 •

edited

Loading