[RL][Frontend] Return token-level weight-version spans on final outputs

[ShuoleiWang]

Purpose

Draft for early testing. This is a first version for others to try; it will change.

RL rollouts often cross a weight update. With pause(mode="keep"), a running request is frozen, the weights and the weight-version label change, and the request resumes, so one output is sampled by two weight versions. vLLM only reports the engine's current label (/weight_info), so the caller cannot tell which tokens came from which weights.

This PR returns, on the final output of every sequence, the label of the weights that sampled each range of output tokens:

"weight_versions": [
  {"version": "v7", "start": 0,  "end": 57},
  {"version": "v8", "start": 57, "end": 128}
]

Behavior:

• Shape. There is one list per output sequence, made of half-open ranges over output-token indexes (prompt excluded). The ranges are contiguous and in generation order, and they cover every returned token once.
• Labels. Adjacent ranges with the same label merge; A → B → A gives three ranges. An output with no tokens gets [].
• Capture point. A token's label is the engine's label when the step that sampled it returns its outputs. The update flow sets the label while generation is paused, when no step is in flight, so every token gets the label of the weights that sampled it.
• Delivery. The list is returned only on the final output, including aborted outputs. A stream sends it on its last chunk, and that chunk is still sent when it has no new tokens (e.g. an abort).
• Surfaces. The list is exposed as CompletionOutput.weight_versions and on /inference/v1/generate choices in both output modes (non-stream and the final stream chunk). OpenAI-compatible endpoints are unchanged in this PR.

Implementation:

• EngineCore stamps its current label on each step's EngineCoreOutputs, next to the existing iteration details. The scheduler and the weight-update path are unchanged.
• OutputProcessor records a span start whenever a request's tokens arrive under a different label than its previous ones, and builds the spans on the final output.
• RequestOutput.add takes the spans from the last merged output, like finish_reason. Merged DELTA outputs therefore keep them, and output merged past the end of a streaming-input chunk does not carry stale ones.
• The Rust engine-core client decodes the new EngineCoreOutputs field, because it rejects unknown trailing fields of that struct. It does not expose the field yet.
• Docs: the token-in/token-out and async RL pages describe the field and where to set the label. The async RL example now reads the swap point from weight_versions instead of pausing, sleeping five seconds and counting streamed tokens. CI runs that example, so its check against a fresh instance of the new model now also covers the span boundary.

Not a duplicate:

• Another open draft reports one weight version per output: each streamed chunk carries the label current when it was formed, and a non-streaming response carries the label of its last output.
• That cannot say which tokens of a non-streaming response came from which weights, and a streamed chunk that merges outputs from both sides of an update carries one label.
• This PR returns the ranges themselves, per completion. Token-level masking and per-segment replay of routed_experts or sampling_mask need these offsets.

This change only adds metadata. It does not change scheduling, sampling or the generated tokens, so no model evaluation is included.

Test Plan

Unit tests (CPU):

pytest tests/v1/engine/test_output_processor.py -k weight_version
cargo test -p vllm-engine-core-client

End-to-end on H20, on this revision: the precompiled wheel of the base commit with this PR's Python files on top.

• /inference/v1/generate. Two servers with Qwen3-1.7B, VLLM_SERVER_DEV_MODE=1 and --enable-scale-out, one with async scheduling and one without.
  • Weight update. A worker extension zeroes lm_head through /collective_rpc, and /update_weight_version sets a new label. After that, greedy decoding emits only token 0, so each span boundary must sit at the first 0 token.
  • Scenarios, in both output modes (tokens and text):
    • pause(mode="keep") with clear_cache=false and with clear_cache=true, using three streams and one non-streaming request, plus a request admitted during the pause;
    • pause(mode="abort") with streams, a non-streaming request and a long-prompt request sent just before the abort;
    • a request after resume, and /weight_info.
• Example. examples/rl/rlhf_async_new_apis.py on two GPUs, with only the model names pointed at local copies of the same checkpoints:

VLLM_ALLOW_INSECURE_SERIALIZATION=1 python examples/rl/rlhf_async_new_apis.py

Test Result

• Unit tests: the new test passes in all 4 parametrizations, and vllm-engine-core-client passes 125/125.
• Existing tests in the touched areas (the output processor, scheduler and outputs tests, and all /inference/v1/generate tests): they give the same results as on the base commit, including the same 59 errors on both, which come from the CPU test environment (e.g. model weights not available offline).
• End-to-end, /inference/v1/generate: 135/135 checks pass with async scheduling on, and 135/135 with it off.
  • keep: streams and non-streaming requests got [("A", 0, k), ("B", k, 300)] with k between 17 and 21. Every token from k on was 0, and the token before k was not. Requests admitted during the pause got [("B", 0, 64)].
  • abort: every aborted request ended with finish_reason: "abort" and one span over its partial output, including final chunks with no new tokens. The long-prompt request was aborted before its first token once (async scheduling on, tokens mode) and got []. In the other three runs it had already produced one token and got one span.
  • after resume: one span with the unchanged label, and /weight_info returned the last label.
• Example: passes, with 13/13 prompts matching a fresh instance of the new model. The swap points read from weight_versions were 10 tokens for 12 prompts and 12 tokens for one.

Coverage notes:

• The unit test covers the span logic in OutputProcessor. EngineCore stamping the label and the /inference/v1/generate changes are covered by the end-to-end runs and the example.
• The streaming-input chunk-end case of RequestOutput.add is not tested.

On current main, /inference/v1/generate is only registered with --enable-scale-out (or --tokens-only).

AI assistance (Claude Code) was used to write the code, the tests and this description.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ShuoleiWang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ShuoleiWang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Documentation preview: https://vllm--59582.org.readthedocs.build/en/59582/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ShuoleiWang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @ShuoleiWang, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[ShuoleiWang]

/ci run

[github-actions]

❌ This PR is 7 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[ShuoleiWang]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94038 for commit c07551e1edbf.