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

docs/serving/online_serving/token_in_token_out.md

@@ -55,6 +55,21 @@ The server returns a 400 for:
 
 For using `output_mode` with separate prefill and decode pools, see [Disaggregated Prefilling](../../features/disagg_prefill.md#generate-api-output-modes).
 
+## Weight versions
+
+The final output of every choice carries `weight_versions`, the label of the weights that sampled each range of its output tokens. Ranges are half-open over output-token indexes (prompt excluded), contiguous and in generation order, so a request that crossed one weight update gets two:
+
+```json
+"weight_versions": [
+  {"version": "v7", "start": 0, "end": 57},
+  {"version": "v8", "start": 57, "end": 128}
+]
+```
+
+- With `stream: true`, only the last chunk of each choice carries the list, including the last chunk after an abort. Its offsets cover the whole output. Earlier chunks carry `null`.
+- An output with no tokens, such as a request aborted before its first token, gets `[]`.
+- The labels are the ones you set with `/update_weight_version` or `/finish_weight_update`, see [Async Reinforcement Learning](../../training/async_rl.md#weight-versions-in-outputs). vLLM doesn't order or check them, and the label is `"default"` until you set one.
+
 ## Aborting requests
 
 `POST /inference/v1/abort_requests` aborts in-flight requests. It is registered wherever `/inference/v1/generate` is and requires the API key when `--api-key` is set, like `/inference/v1/generate`.

docs/training/async_rl.md

@@ -43,23 +43,39 @@ With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionali
 - `POST /pause?mode=keep` - Pause generation
 - `POST /resume` - Resume generation
 - `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`)
+- `POST /update_weight_version` - Set the `weight_version` label without changing weights (send `{"new_version": "v8"}`)
 - `GET /weight_info` - Return the latest committed `weight_version`
 
 !!! note "Data Parallelism"
-    When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.
+    When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update, and set the weight version on every instance too.
 
 ## Typical Async RL Flow
 
 A typical async RL loop with weight syncing looks like this:
 
 1. Start generating rollouts from the current policy
 2. Once trainer has new weights to update to, pause generation with `mode="keep"`
-3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md))
+3. Sync the updated weights from the trainer to the inference engine (see [Weight Transfer](weight_transfer/README.md)) and set their weight version
 4. Resume generation -- in-flight requests continue with the new weights
 5. Repeat
 
 The key insight is that requests paused with `mode="keep"` will produce tokens from the **old** weights before the pause and tokens from the **new** weights after resume. The `clear_cache` parameter controls whether the KV cache is invalidated during the pause. When `clear_cache=True`, previously cached key-value entries are discarded, so all tokens generated after resume will be computed entirely with the new weights. When `clear_cache=False`, existing KV cache entries are retained, meaning some tokens in context may still reflect the old weights (stale KV cache).
 
+## Weight Versions in Outputs
+
+To tell which tokens came from which weights, set a weight-version label while generation is paused, after the new weights are loaded: pass `weight_version` to `finish_weight_update`, or call `update_weight_version` (`POST /update_weight_version` over HTTP). Labels are strings you choose, such as the trainer step; vLLM doesn't order or check them.
+
+The final output of every sequence then carries `weight_versions`, one span per label over its output-token indexes. A request paused with `mode="keep"` across one update gets two spans:
+
+```python
+completion = final_output.outputs[0]
+for span in completion.weight_versions:
+    # The tokens in [span.start, span.end) were sampled by weights `span.version`.
+    print(span.version, completion.token_ids[span.start : span.end])
+```
+
+The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. `/inference/v1/generate` returns the same list on each choice; see [Tokens In <> Tokens Out API](../serving/online_serving/token_in_token_out.md#weight-versions).
+
 ## Example
 
 The [async RLHF example](../../examples/rl/rlhf_async_new_apis.py) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation.

examples/rl/rlhf_async_new_apis.py

@@ -26,8 +26,8 @@
 * Pause generation once any request reaches a token threshold.
 * Broadcast the training model's weights to the inference engine
   via the NCCL weight transfer engine, replacing the base weights.
-* Resume generation and collect results, noting which tokens were
-  generated before vs. after the weight swap.
+* Label the new weights, resume generation, and use `weight_versions`
+  to locate the weight swap in each output.
 * Validate correctness by launching a fresh vLLM instance loaded
   directly with the training model and comparing its output to the
   post-swap tokens from the weight-synced engine.
@@ -61,6 +61,7 @@
 MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
 MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
 PAUSE_TOKEN_THRESHOLD = 10
+WEIGHT_VERSION_V2 = "v2"
 ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"
 
 
@@ -77,44 +78,27 @@ def __init__(self, **kwargs):
             log_requests=engine_args.enable_log_requests,
             log_stats=not engine_args.disable_log_stats,
         )
-        self._generation_paused = False
         self._request_pause_flag = False
 
     async def do_generate(
         self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
-    ) -> tuple[vllm.RequestOutput, int]:
-        """Generate a single request, setting the request pause flag once the
-        token count reaches the threshold.
-
-        Returns (output, pause_token_index). pause_token_index is the number
-        of tokens generated before the weight change, or -1 if no pause.
-        """
-        pause_token_index = -1
-        prev_token_count = 0
+    ) -> vllm.RequestOutput:
+        """Generate a request and set the pause flag at the token threshold."""
         async for request_output in self.generate(
             {"prompt_token_ids": prompt_token_ids},
             sampling_params,
             request_id=str(uuid.uuid4()),
         ):
             output = request_output
-            cur_token_count = len(output.outputs[0].token_ids)
-            if (
-                cur_token_count >= PAUSE_TOKEN_THRESHOLD
-                and not self._request_pause_flag
-            ):
+            if len(output.outputs[0].token_ids) >= PAUSE_TOKEN_THRESHOLD:
                 self._request_pause_flag = True
-            if self._generation_paused and pause_token_index == -1:
-                pause_token_index = prev_token_count
-            prev_token_count = cur_token_count
-        return output, pause_token_index
+        return output
 
     async def pause_after_n_tokens(self):
         """Wait for any request to set the pause flag, then pause."""
         while not self._request_pause_flag:
             await asyncio.sleep(0)
         await super().pause_generation(mode="keep")
-        await asyncio.sleep(5)
-        self._generation_paused = True
 
 
 @ray.remote(num_gpus=1)
@@ -262,9 +246,18 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
 ray.get(llm.pause_after_n_tokens.remote())
 
 ray.get(train_model.broadcast_weights.remote())
+# Label the new weights before resuming generation.
+ray.get(llm.update_weight_version.remote(WEIGHT_VERSION_V2))
 
 ray.get(llm.resume_generation.remote())
-results = ray.get(gen_futures)
+results = []
+for output in ray.get(gen_futures):
+    completion = output.outputs[0]
+    pause_idx = next(
+        (s.start for s in completion.weight_versions if s.version == WEIGHT_VERSION_V2),
+        len(completion.token_ids),
+    )
+    results.append((output, pause_idx))
 
 for i, (output, pause_idx) in enumerate(results):
     all_token_ids = list(output.outputs[0].token_ids)
@@ -326,7 +319,7 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
 
 num_pass = 0
 num_total = len(results)
-for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
+for i, ((output, pause_idx), val_output) in enumerate(zip(results, val_results)):
     expected = list(output.outputs[0].token_ids)[pause_idx:]
     actual = list(val_output.outputs[0].token_ids)
     match = actual == expected

rust/src/engine-core-client/src/protocol/output.rs

@@ -228,6 +228,9 @@ struct WireEngineCoreOutputs {
     /// wave needs to start in other engines.
     #[serde(default)]
     start_wave: Option<u32>,
+    /// Decoded for Python wire compatibility; not exposed by the Rust frontend.
+    #[serde(default)]
+    weight_version: Option<String>,
 }
 
 /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`.

rust/src/engine-core-client/src/tests/python_compat.py

@@ -139,6 +139,7 @@ class EngineCoreOutputs(
     finished_requests: set[str] | None = None
     wave_complete: int | None = None
     start_wave: int | None = None
+    weight_version: str | None = None
 
 
 request = EngineCoreRequest(
@@ -240,6 +241,7 @@ class EngineCoreOutputs(
         )
     ],
     finished_requests={"req-1"},
+    weight_version="v7",
 )
 
 extended_outputs = EngineCoreOutputs(
@@ -252,6 +254,7 @@ class EngineCoreOutputs(
         )
     ],
     finished_requests=outputs.finished_requests,
+    weight_version=outputs.weight_version,
 )
 extended_outputs_bytes = msgspec.msgpack.encode(extended_outputs)
 # The ordinary frontend's schema ignores even non-default extension values.

tests/v1/engine/test_output_processor.py

@@ -24,6 +24,7 @@
     PoolingRequestOutput,
     RequestOutput,
     SamplingMask,
+    WeightVersionSpan,
 )
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tokenizers import TokenizerLike
@@ -87,6 +88,7 @@ def test_completion_output_preserves_each_sampling_mask_position() -> None:
     ]
     state.routed_experts_chunks = []
     state.spec_decode_metrics = None
+    state.weight_version_starts = []
 
     output = state._new_completion_output([1, 2, 3], FinishReason.LENGTH, None)
 
@@ -1596,6 +1598,7 @@ def test_sampling_masks_follow_output_kind(output_kind):
     state.sampling_mask_chunks = []
     state.routed_experts_chunks = []
     state.spec_decode_metrics = None
+    state.weight_version_starts = []
 
     supports = [[10, 11], [20], [30, 31]]
     masks = []
@@ -1725,3 +1728,53 @@ def test_async_pooling_cache_miss_does_not_fail_output_processor():
     assert output.error is not None
     assert output.error.retryable
     assert not output_processor.has_unfinished_requests()
+
+
+@pytest.mark.parametrize(
+    "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY]
+)
+@pytest.mark.parametrize(
+    "steps,expected",
+    [
+        # Merge adjacent equal labels; A -> B -> A remains three spans.
+        (
+            [("v7", [10, 11]), ("v7", [12]), ("v8", [13, 14]), ("v7", [15])],
+            [("v7", 0, 3), ("v8", 3, 5), ("v7", 5, 6)],
+        ),
+        # An output without tokens has no spans.
+        ([], []),
+    ],
+)
+def test_weight_version_spans(output_kind: RequestOutputKind, steps, expected):
+    output_processor = OutputProcessor(tokenizer=None, log_stats=False)
+    queue = RequestOutputCollector(output_kind, request_id="request-0-int")
+    request = EngineCoreRequest(
+        request_id="request-0-int",
+        external_req_id="request-0",
+        prompt_token_ids=[1, 2, 3],
+        mm_features=None,
+        arrival_time=0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        sampling_params=SamplingParams(detokenize=False, output_kind=output_kind),
+        pooling_params=None,
+    )
+    output_processor.add_request(request, prompt=None, queue=queue)
+
+    for i, (weight_version, token_ids) in enumerate(steps):
+        # Copy: merging deltas extends token lists in place.
+        output = EngineCoreOutput(
+            request_id="request-0-int", new_token_ids=list(token_ids)
+        )
+        output_processor.process_outputs([output], weight_version=weight_version)
+        if i == 0 and output_kind == RequestOutputKind.DELTA:
+            intermediate = queue.get_nowait()
+            assert intermediate is not None
+            assert intermediate.outputs[0].weight_versions is None
+
+    # Unread deltas merge into the abort output without losing spans.
+    output_processor.abort_requests(["request-0"], internal=False)
+    final = queue.get_nowait()
+    assert final is not None and final.finished
+    assert final.outputs[0].weight_versions == [WeightVersionSpan(*s) for s in expected]

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

@@ -31,6 +31,7 @@
 )
 from vllm.entrypoints.serve.engine.protocol import UsageInfo
 from vllm.logprobs import Logprob
+from vllm.outputs import WeightVersionSpan
 from vllm.renderers import TokenizeParams
 from vllm.sampling_params import SamplingParams
 from vllm.utils import random_uuid
@@ -411,6 +412,8 @@ class GenerateChoiceBase(BaseModel):
     # or (b) `enable_return_routed_experts` is off server-side.
     routed_experts: str | None = None
     sampling_mask: list[list[int]] | None = None
+    # Contiguous weight version spans over output-token indexes.
+    weight_versions: list[WeightVersionSpan] | None = None
 
     @field_validator("token_ids")
     @classmethod
@@ -443,6 +446,8 @@ class GenerateStreamChoiceBase(BaseModel):
     token_ids: list[int] | None = None
     routed_experts: str | None = None
     sampling_mask: list[list[int]] | None = None
+    # Set on the final chunk only; covers the whole output.
+    weight_versions: list[WeightVersionSpan] | None = None
 
 
 class GenerateTokensStreamChoice(GenerateStreamChoiceBase):

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -453,6 +453,7 @@ async def serve_tokens_full_generator(
                 token_ids=as_list(output.token_ids),
                 routed_experts=routed_experts_b64,
                 sampling_mask=sampling_mask,
+                weight_versions=output.weight_versions,
             )
             if text_mode:
                 text_choices.append(
@@ -625,6 +626,7 @@ async def serve_tokens_stream_generator(
                         token_ids=as_list(delta_token_ids),
                         routed_experts=routed_experts_b64,
                         sampling_mask=sampling_mask,
+                        weight_versions=output.weight_versions,
                     )
                     chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse
                     if text_mode:

vllm/outputs.py

@@ -38,6 +38,15 @@ class SamplingMask:
     token_ids: list[list[int]]
 
 
+@dataclass
+class WeightVersionSpan:
+    """Output tokens `[start, end)` were sampled by weights labelled `version`."""
+
+    version: str
+    start: int
+    end: int
+
+
 @dataclass
 class CompletionOutput:
     """The output data of one completion output of a request.
@@ -62,6 +71,8 @@ class CompletionOutput:
             ``--per-request-spec-decode-metrics`` is enabled; None otherwise.
             Surfaced in the response as ``metrics.speculative_decoding`` for
             single-sequence (``n == 1``) requests.
+        weight_versions: Contiguous weight version spans over output-token indexes.
+            Set on the final output only.
 
     """
 
@@ -76,6 +87,7 @@ class CompletionOutput:
     lora_request: LoRARequest | None = None
     sampling_mask: SamplingMask | None = None
     spec_decode_metrics: RequestSpecDecodeMetrics | None = None
+    weight_versions: list[WeightVersionSpan] | None = None
 
     def finished(self) -> bool:
         return self.finish_reason is not None
@@ -223,6 +235,8 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
                         # aggregation with earlier chunks that have no R3.
                         if next_completion.routed_experts is not None:
                             completion.routed_experts = next_completion.routed_experts
+                        # Like finish_reason, spans come only with the terminal output.
+                        completion.weight_versions = next_completion.weight_versions
                         completion.finish_reason = next_completion.finish_reason
                         completion.stop_reason = next_completion.stop_reason
                     else:

vllm/v1/engine/__init__.py

@@ -282,6 +282,9 @@ class EngineCoreOutputs(
     # "old" wave, so the next wave needs to be started in other engines.
     start_wave: int | None = None
 
+    # Weight version used for this step.
+    weight_version: str | None = None
+
     def __post_init__(self):
         if self.timestamp == 0.0:
             self.timestamp = time.monotonic()

vllm/v1/engine/async_llm.py

@@ -855,7 +855,10 @@ async def output_handler():
                         outputs_slice = engine_core_outputs[start:end]
                         # 2) Process EngineCoreOutputs.
                         processed_outputs = output_processor.process_outputs(
-                            outputs_slice, outputs.timestamp, iteration_stats
+                            outputs_slice,
+                            outputs.timestamp,
+                            iteration_stats,
+                            weight_version=outputs.weight_version,
                         )
                         # NOTE: RequestOutputs are pushed to their queues.
                         assert not processed_outputs.request_outputs