[Feat] RL: Add radix transfer plugins by guapisolo · Pull Request #42647 · sgl-project/sglang

guapisolo · 2026-10-05T20:12:42Z

Depends on: #42646
Stacked on the parent PR's branch. Review / merge the parent first.

Summary

Return replay outputs through a shared output-store package instead of inline.

Motivation

Routed experts, indexer top-k, and sampling masks dominate long RL rollout responses as inline base64 or lists: 6.3 MB for 2048 prompt and 1024 output tokens. Readers in other processes need the same field names, ref format, and storage client.

Usage

python -m sglang.launch_server --model-path <model> --enable-return-routed-experts \
  --output-store-backend mooncake --output-store-backend-extra-config \
  '{"master_server_address": "10.0.0.1:50051", "local_hostname": "10.0.0.2",
    "local_buffer_size": "1gb", "key_prefix": "run-1/output-store"}'
curl -s localhost:30000/generate -d '{"text": "Hi", "return_routed_experts": true,
  "return_outputs_via_store": true}' | jq '.meta_info.output_store_ref'

# pip install -e radix_transfer_plugins: needed by readers, and by servers only with the store
store = OutputStore(backend=create_backend("mooncake", config_json))
ref = StoreRef.from_builtins(meta_info["output_store_ref"])
# Many objects, chunked to the client's read buffer and copied out: arrays the caller owns.
(stash,) = store.read_many([(ref, ["routed_experts"])], stash_type=TokenOutputStash, buffer_bytes=1 << 30)
store.delete_many([ref])

Design Notes

Mental model: an opted-in request collects its outputs in a TokenOutputCollector; on a successful finish OutputStoreWriter stores them with one OutputStore.put_many off the event loop, so meta_info carries a StoreRef instead of the arrays; the package owns the stored format, SGLang owns collection plus cleanup, and other requests stay inline.
radix_transfer_plugins imports neither sglang nor torch; SGLang imports it only with --output-store-backend. read_many copies chunks out because Mooncake views pin the read buffer.
Cleanup: the writer deletes objects whose ref is never sent (cancellation, disconnect, failed chat response), not siblings of a failed batch. Mooncake removes a failed put_batch's keys and retries a put that found the local buffer full once, alone.
Only support-sized mask rows are stored; sampled tokens' logprobs stay in meta_info (support mode: output_token_sampling_selected_logprobs).

Verification

radix_transfer_plugins/tests (macOS, fake Mooncake): 32 passed, including two puts that both find the buffer full (hang before this head).
test_output_store.py, test_tokenizer_manager_rid_cleanup.py, test_sampling_mask_validation.py, test_multi_tokenizer_mixin.py, test_batch_result_processor_hidden_states.py, test_output_streamer_customized_info.py, test_serving_chat.py (macOS CPU): 279 passed, including a server import without the package and a zero-output prefill (both fail before this head).
test_server_args_docs_coverage.py fails only on 8 base-branch flags.
Not run: GPU and a real Mooncake client.

Review Focus

TokenizerManager._store_token_outputs: cancellation and disconnect cleanup.
add_sampling_mask_return_values / SamplingMaskChunk.selected_logprobs: one mask row per returned token.
_LocalBufferGate: a waiting put holds no buffer.

CI States

Latest PR Test (Base): ❌ Run #39680011874
Latest PR Test (Extra): ❌ Run #39680011858
Latest PR Test (AMD ROCm 10): ❌ Run #39680011757

Even as segments, routed experts, indexer top-k and the sampling mask still travel inline in every SGLang response and through the rollout data pack. With --use-output-store, the engines write them to the run's Mooncake through sgl-project/sglang#42647 and return a small output_store_ref, and the trainer reads them. - miles/utils/output_store_config.py owns --use-output-store and --output-store-local-buffer-size, the launch checks (Mooncake object store, no PD, in_place pause, Ray only, no train fault tolerance, no multi-policy, no external engines) and the engines' output store config; per-group overrides of the SGLang store flags are rejected. - Training requests that ask for replay outputs set return_outputs_via_store; evaluation stays inline. ReplayIndex parses the ref into store segments, checks structure from the ref's field specs without reading, keeps one source per field of a sample, and rejects a store request answered inline (a router may have dropped the flag) or a ref of another backend. - The trainer reads each (key, field) once, in chunks that fit the read buffer, copies them out, and fails before reading when one object is larger than the buffer. - Store refs go to the train dict as plain JSON types, which the Mooncake object store's strict msgpack accepts.

Even as segments, routed experts, indexer top-k and the sampling mask still travel inline in every SGLang response and through the rollout data pack. With --use-output-store, the engines write them to the run's Mooncake through sgl-project/sglang#42647 and return a small output_store_ref, and the trainer reads them. - miles/utils/output_store_config.py owns --use-output-store and --output-store-local-buffer-size, the launch checks (Mooncake object store, no PD, in_place pause, Ray only, no train fault tolerance, no multi-policy, no external engines) and the engines' output store config; per-group overrides of the SGLang store flags are rejected. - Training requests that ask for replay outputs set return_outputs_via_store; evaluation stays inline. ReplayIndex parses the ref into store segments, checks structure from the ref's field specs without reading, keeps one source per field of a sample, and rejects a store request answered inline (a router may have dropped the flag) or a ref of another backend. - The trainer reads each (key, field) once, in chunks that fit the read buffer, copies them out, and fails before reading when one object is larger than the buffer. - Store refs go to the train dict as plain JSON types, which the Mooncake object store's strict msgpack accepts. - --use-output-store turns on --enable-delay-token-replay-assembly at launch: the rollout executor cannot read store objects, so the train ranks assemble the replay data. - A store response carries no support rows inline, so in support mode its rollout_log_probs are SGLang's output_token_sampling_selected_logprobs; inline responses still pick each sampled token's entry from its row.

…store A request that sets return_outputs_via_store keeps its routed experts, indexer top-k and sampling mask out of the response: the tokenizer stores them when the request finishes and returns meta_info["output_store_ref"] instead. The storage layer is a new standalone package at the repository root, radix_transfer_plugins (msgspec and numpy only, no sglang import), so readers such as RL trainers share the stash schema, ref format, codecs and backends with SGLang: - StoreRef is self-describing (backend, full key, and each field's dtype, shape, codec and size); OutputStore writes one response with one put_many, all or nothing. - The Mooncake backend stores one hard-pinned key per field through the native store API, removes every key of a failed put_batch, removes with force because a read leaves a lease, and retries a put that finds the client's local buffer full once another put frees room. On the SGLang side, output_store.py keeps the per-request collector and a writer that stores off the event loop and deletes the objects whose ref is never sent (cancellation, client disconnect, chat response build failure). --output-store-backend accepts any backend registered in the package. This replaces the MooncakeBundleTransfer-based store and its ref format.

A request that returned its outputs via the output store also stored the per-token sampling logprobs, so in selected mode the client got no sampled-token logprob in the response and had to read the store for one float per token. Only support-sized data now goes to the store: - selected mode keeps output_token_sampling_logprobs in meta_info, and the stash no longer holds it; - support mode also returns output_token_sampling_selected_logprobs, each sampled token's logprob picked from its support by output id, with or without the store; the stash keeps the support logprobs. The pick runs in the tokenizer, so the scheduler, the IPC structs and the PD transfer are unchanged.

A reader that needs many objects (an RL trainer reading every object of its data shard) could only call get_many, whose arrays are views over the backend's read buffers. With Mooncake those views hold the client's local buffer, so a reader had to chunk its reads and copy each chunk out before the next one itself, and a view kept alive past its chunk overflowed the buffer on the next read. read_many(requests, stash_type=, buffer_bytes=) takes (ref, field names) pairs and owns that lifecycle: it fails before reading anything when a field is missing or larger than the buffer, groups the fields in order into chunks whose stored bytes fit buffer_bytes, and copies each chunk into new memory before reading the next, so the backend holds at most one chunk. It knows nothing about what the objects mean.

…_len, retry a full buffer alone - Import radix_transfer_plugins only when --output-store-backend is set, so servers and tests that never enable the store do not need the package; the CPU stage installs it for the output store tests. - A sampled prefill that finishes with no output tokens (max_new_tokens=0, as the n>1 prefix warmup sends) no longer returns a sampling mask row; in support mode SamplingMaskChunk.selected_logprobs asserted on it and the tokenizer process exited. - _LocalBufferGate: concurrent puts that all found the local buffer full waited on each other forever. A failed put now retries once with no other put running, and its status is final. - Document --output-store-backend and --output-store-backend-extra-config; comment fixes.

guapisolo · 2026-10-21T22:11:38Z

Need discussion: Where to maintain SoT?

guapisolo requested review from JustinTong0323, Ying1123, hnyls2002, merrymercy and xiezhq-hermann as code owners October 5, 2026 20:12

guapisolo closed this Oct 5, 2026

guapisolo reopened this Oct 6, 2026

This was referenced Oct 15, 2026

Resolve SGLang output-store references in rollout generation radixark/miles#3907

Closed

feat(protocol): forward return_outputs_via_store on generate and chat radixark/sgl-router-for-miles#23

Merged