Skip to content

feat: support streaming external rollouts - #2272

Open
Aphoh wants to merge 1 commit into
THUDM:mainfrom
Aphoh:warnold/external-streaming-generate
Open

feat: support streaming external rollouts#2272
Aphoh wants to merge 1 commit into
THUDM:mainfrom
Aphoh:warnold/external-streaming-generate

Conversation

@Aphoh

@Aphoh Aphoh commented Aug 14, 2026

Copy link
Copy Markdown

Summary

This PR makes Slime's external SGLang rollout path work with streaming /generate responses.

It adds:

  • cumulative and incremental SSE stream handling
  • request-level abort for streaming generators, so partial rollouts cancel only their active HTTP requests and retain observed prefixes
  • incremental accumulation of tokens, logprobs, top-p metadata, routed-expert metadata, and response text
  • validation that rejects malformed or mismatched stream formats instead of silently training on incomplete output

No dependencies are added.

Why

The non-streaming rollout path receives generation state only after /generate completes. If a rollout is interrupted early, it cannot retain the generated prefix without a separate polling or abort-response API.

Streaming exposes that prefix as it is generated. Slime can apply each observed chunk directly to the Sample, cancel the request when partial-rollout collection fires, and later resume from the retained tokens.

Incremental streaming is also important for long contexts: Slime processes only the new suffix from each chunk and materializes response text once, avoiding repeated concatenation or rescanning of the accumulated response. Cumulative streams remain supported for compatibility.

Design

Stream normalization

SGLangStreamAccumulator normalizes SGLang's two stream formats into updates that can be applied directly to a sample:

  • cumulative: each chunk repeats the response so far; Slime validates the reported cumulative length and applies only the unseen suffix
  • incremental: each chunk contains only new output; Slime validates that the prior length plus the chunk length equals SGLang's reported length

The accumulator handles empty terminal chunks, incremental or terminal-full top-p snapshots, encoded metadata arrays, routed-expert metadata, Unicode and special-token text, and resumed samples. Text chunks are collected and joined once; when text is absent, accumulated token IDs are decoded once.

The stream must end with a finish_reason. An unexpected EOF raises an error rather than allowing an incomplete sample into reward computation or training.

Abort lifecycle

A generation function opts into request-level abort by setting an attribute:

generate_streaming.abort_mode = "request"

Slime tracks active request-abortable tasks separately from active server-abort generations:

  • request-abortable tasks are canceled individually
  • server-wide SGLang abort is used only while a server-abort generator is active
  • CancelledError is converted into an ABORTED sample only when Slime deliberately canceled that task; unrelated cancellation still propagates
  • custom generation functions without the attribute retain the existing server-wide abort behavior

This keeps the lifecycle generic: custom generation functions can opt in without adding function paths to a registry or introducing a global abort-mode CLI option.

Partial-rollout resumption

Slime already skips completed or truncated samples when a buffered group is reused. The request-abort path preserves that behavior: a group is buffered only when at least one aborted sample has a non-empty observed prefix, completed siblings remain unchanged, and only aborted siblings resume generation. Existing off-policy masking and reward behavior are also preserved.

Configuration

Select the streaming generator:

--custom-generate-function-path slime.rollout.sglang_streaming_rollout.generate_streaming

For cumulative SGLang output, no additional Slime flag is needed.

For incremental output, configure both sides:

# SGLang
--incremental-streaming-output

# Slime
--sglang-incremental-streaming-output

The Slime flag must match the server's output mode. Length validation fails early if cumulative output is sent to the incremental consumer or vice versa.

Request cancellation can preserve only metadata received before disconnect. Metadata that SGLang emits only in a terminal chunk cannot be recovered from a canceled partial request; this PR does not attempt to reconstruct it.

Diff size

Measured against the current PR base:

Area Files Added Removed Net
Production code 4 265 80 +185
Tests 2 529 0 +529
CI registration 2 5 0 +5
Total 8 799 80 +719

The test contribution is 527 lines in the new focused streaming suite plus two fixture lines adapting the existing custom-generator contract tests. CI changes only register the new CPU test file in the generated and checked-in matrices.

Automated validation

The new streaming suite contains 26 parameterized cases covering:

  • cumulative and incremental streams
  • stream intervals 1, 20, and 64
  • incremental and terminal-full top-p metadata
  • encoded routed-expert and top-p arrays
  • Unicode, special-token text, absent text, and empty terminal chunks
  • stop and length finish reasons
  • inconsistent lengths, mismatched stream modes, and unexpected EOF
  • request cancellation, unrelated parent cancellation, and server abort behavior
  • mixed completed/aborted sibling groups through abort, buffer insertion, and resume without duplicated tokens or rewards

The five existing custom-generator contract tests and fourteen argument-validation regression tests were also run, for 45 focused passing cases in total.

Additional checks:

  • direct python tests/... entrypoints passed for the modified test files
  • uvx prek run --all-files passed
  • a synthetic incremental accumulator benchmark processed 500,000 output tokens in 25,000 chunks at stream interval 20 in 0.132 seconds median across five runs

Live external-SGLang E2E training

The PR at 4d39b5a606a410c80dd74475a2467b9a94a137fa was built into one image and that exact image was used for both the trainer and external SGLang engine in both runs.

Shared setup:

  • Qwen3-30B-A3B, exercising a 128-expert MoE model
  • vanilla external SGLang /generate
  • SGLang TP4/EP4 inference and Megatron TP4/EP4 training on disjoint four-GPU placements
  • SGLang --stream-interval 20
  • SGLang --enable-return-routed-experts
  • Slime --use-rollout-routing-replay
  • two prompts with two samples per prompt
  • 256 maximum response tokens, one rollout, and global training batch size 4
  • a real Megatron forward/backward/optimizer step
  • full trainer-to-inference NCCL weight synchronization before generation and after the optimizer step, transferring all 116 model-weight buckets each time

The only A/B difference was stream representation:

Run SGLang Slime
Cumulative default cumulative streaming default cumulative handling
Incremental --incremental-streaming-output --sglang-incremental-streaming-output

Results:

Result Cumulative Incremental
Completed responses 4/4 4/4
Response tokens per sample 256 256
Rollout generation time 26.1 s 26.3 s
Megatron optimizer step passed passed
Gradient norm 0.04534 0.03906
Post-train weight update 116/116 buckets 116/116 buckets

All responses reached the configured length limit, exercising streamed length termination. Saved rollout artifacts were inspected directly. Every sample in both modes had:

  • 285 total tokens: 29 prompt tokens and 256 response tokens
  • rollout_routed_experts.dtype == torch.int32
  • rollout_routed_experts.shape == (284, 48, 8), matching (len(tokens) - 1, num_moe_layers, experts_per_token)
  • valid expert IDs from 0 through 127
  • routed-expert data in all 48 MoE layers

Both modes consumed those tensors with routing replay enabled and completed forward, backward, optimizer, and post-train weight synchronization.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant