[RL][Frontend] Return weight-version spans on OpenAI completions and chat (#61347)

Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>

docs/training/async_rl.md

@@ -76,6 +76,10 @@ for span in completion.weight_versions:
 
 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).
 
+`/v1/completions` and `/v1/chat/completions` return it on each choice when the request sets `return_weight_versions: true`. The offsets count the choice's output tokens, which `return_token_ids` returns as `token_ids`, and a streamed choice gets the list on its last chunk.
+
+An aborted request returns spans over its partial output. If you resubmit that output as part of a new prompt, the new request's spans start again at 0, so shift them by the number of tokens already generated before joining the two lists.
+
 ## 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.

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -1188,6 +1188,15 @@ async def test_streaming_reasoning_usage_lifecycle_edges(
         await iterator.aclose()
 
     error_chunks = [c for c in chunks if isinstance(c, dict) and "error" in c]
+    if name == "empty-finish":
+        # A choice that finishes without output still gets its last chunk, unparsed.
+        assert any(
+            choice.get("finish_reason") == "stop"
+            for item in chunks
+            if isinstance(item, dict)
+            for choice in item.get("choices", [])
+        )
+        parsers[0].parse_delta.assert_not_called()
     call_count = sum(parser.count_reasoning_tokens.call_count for parser in parsers)
     if expected_tokens is not None:
         assert outcome == "exhausted" and chunks[-1] == "[DONE]" and not error_chunks

tests/entrypoints/openai/completion/test_completion_error.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from dataclasses import dataclass, field
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
@@ -17,7 +18,7 @@ from vllm.entrypoints.scale_out.render.serving import ServingRender
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.exceptions import GenerationError, VLLMValidationError
 from vllm.logprobs import Logprob
-from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.outputs import CompletionOutput, RequestOutput, WeightVersionSpan
 from vllm.renderers.deepseek_v4 import DeepseekV4Renderer
 from vllm.renderers.hf import HfRenderer
 from vllm.renderers.online_renderer import OnlineRenderer
@@ -310,6 +311,76 @@ def test_completion_metrics_carries_both_timing_and_spec_decode():
     assert response.metrics.speculative_decoding.num_spec_steps == 2
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("opt_in", [False, True])
+async def test_completion_weight_versions(stream: bool, opt_in: bool):
+    """Each choice gets its spans on its final output, even with no tokens."""
+    expected = [{"version": "v1", "start": 0, "end": 2}]
+    spans = [WeightVersionSpan(**span) for span in expected]
+
+    def output(index, token_ids, finish_reason=None, weight_versions=None):
+        return CompletionOutput(
+            index=index,
+            text="x" * len(token_ids),
+            token_ids=token_ids,
+            cumulative_logprob=None,
+            logprobs=None,
+            finish_reason=finish_reason,
+            weight_versions=weight_versions,
+        )
+
+    # Choice 1 is aborted before its first token.
+    final = [output(1, [], "abort", [])]
+    if stream:
+        steps = [[output(0, [100])], [output(0, [101], "length", spans), *final]]
+    else:
+        steps = [[output(0, [100, 101], "length", spans), *final]]
+
+    async def mock_generate(*args, **kwargs):
+        for i, outputs in enumerate(steps):
+            yield RequestOutput(
+                request_id="test-id",
+                prompt="Test prompt",
+                prompt_token_ids=[1, 2, 3],
+                prompt_logprobs=None,
+                outputs=outputs,
+                finished=i == len(steps) - 1,
+            )
+
+    mock_engine = MagicMock(spec=AsyncLLM)
+    mock_engine.errored = False
+    mock_engine.model_config = MockModelConfig()
+    mock_engine.input_processor = MagicMock()
+    mock_engine.renderer = _build_renderer(mock_engine.model_config)
+    mock_engine.generate = MagicMock(side_effect=mock_generate)
+    request = CompletionRequest(
+        model=MODEL_NAME,
+        prompt="Test prompt",
+        max_tokens=10,
+        n=2,
+        stream=stream,
+        return_weight_versions=opt_in,
+    )
+    response = await _build_serving_completion(mock_engine).create_completion(request)
+
+    if stream:
+        finals = {}
+        async for line in response:
+            payload = line.removeprefix("data: ").strip()
+            if payload == "[DONE]":
+                continue
+            for choice in json.loads(payload)["choices"]:
+                if choice.get("finish_reason"):
+                    finals[choice["index"]] = choice
+                else:
+                    assert choice.get("weight_versions") is None
+        got = [finals[i].get("weight_versions") for i in (0, 1)]
+    else:
+        got = [choice["weight_versions"] for choice in response.model_dump()["choices"]]
+    assert got == ([expected, []] if opt_in else [None, None])
+
+
 def _build_online_renderer_for_completion(
     model_config: MockModelConfig,
     renderer: Any | None = None,

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -40,6 +40,7 @@ from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
+from vllm.outputs import WeightVersionSpan
 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
 from vllm.renderers.chat_utils import (
     ChatCompletionMessageParam,
@@ -129,6 +130,8 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
     # ``None`` if (a) the request was aborted before any forward pass,
     # or (b) ``enable_return_routed_experts`` is off server-side.
     routed_experts: str | None = None
+    # Contiguous weight version spans over output-token indexes.
+    weight_versions: list[WeightVersionSpan] | None = None
 
 
 class ChatCompletionResponse(OpenAIBaseModel):
@@ -168,6 +171,8 @@ class ChatCompletionResponseStreamChoice(OpenAIBaseModel):
     stop_reason: int | str | None = None
     # not part of the OpenAI spec but for tracing the tokens
     token_ids: list[int] | None = None
+    # Set on the final chunk only; covers the whole output.
+    weight_versions: list[WeightVersionSpan] | None = None
 
 
 class ChatCompletionStreamResponse(OpenAIBaseModel):
@@ -430,6 +435,15 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "need to map generated text back to input tokens."
         ),
     )
+    return_weight_versions: bool = Field(
+        default=False,
+        description=(
+            "If true, each choice includes weight_versions: the weight-version "
+            "label that sampled each half-open range of its output tokens. "
+            "When streaming, only the last chunk of each choice includes it. "
+            "Beam search does not return it."
+        ),
+    )
     routed_experts_prompt_start: int = Field(
         default=0,
         ge=0,

vllm/entrypoints/openai/chat_completion/serving.py

@@ -651,17 +651,21 @@ class OpenAIServingChat(GenerateBaseServing):
 
                     delta_text = output.text
 
-                    if (
+                    nothing_generated = (
                         not delta_text
                         and not output.token_ids
                         and not previous_num_tokens[i]
-                    ):
+                    )
+                    if nothing_generated and output.finish_reason is None:
                         # Chunked prefill case, don't return empty chunks
                         continue
 
                     delta_message: DeltaMessage | None
 
-                    if parser is not None:
+                    if nothing_generated:
+                        # Finished without output: send the final chunk unparsed.
+                        delta_message = DeltaMessage()
+                    elif parser is not None:
                         delta_message = parser.parse_delta(
                             delta_text=delta_text,
                             delta_token_ids=as_list(output.token_ids),
@@ -773,6 +777,11 @@ class OpenAIServingChat(GenerateBaseServing):
                             token_ids=(
                                 as_list(output.token_ids) if include_token_ids else None
                             ),
+                            weight_versions=(
+                                output.weight_versions
+                                if request.return_weight_versions
+                                else None
+                            ),
                         )
 
                         finish_reason_sent[i] = True
@@ -1097,6 +1106,9 @@ class OpenAIServingChat(GenerateBaseServing):
                     else None
                 ),
                 routed_experts=routed_experts_b64,
+                weight_versions=(
+                    output.weight_versions if request.return_weight_versions else None
+                ),
             )
             choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
 

vllm/entrypoints/openai/completion/protocol.py

@@ -26,6 +26,7 @@ from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
+from vllm.outputs import WeightVersionSpan
 from vllm.renderers import TokenizeParams
 from vllm.sampling_params import (
     BeamSearchParams,
@@ -178,6 +179,15 @@ class CompletionRequest(OpenAIBaseModel):
             "need to map generated text back to input tokens."
         ),
     )
+    return_weight_versions: bool = Field(
+        default=False,
+        description=(
+            "If true, each choice includes weight_versions: the weight-version "
+            "label that sampled each half-open range of its output tokens. "
+            "When streaming, only the last chunk of each choice includes it. "
+            "Beam search does not return it."
+        ),
+    )
     routed_experts_prompt_start: int = Field(
         default=0,
         ge=0,
@@ -677,6 +687,8 @@ class CompletionResponseChoice(OpenAIBaseModel):
     # ``None`` if (a) the request was aborted before any forward pass,
     # or (b) ``enable_return_routed_experts`` is off server-side.
     routed_experts: str | None = None
+    # Contiguous weight version spans over output-token indexes.
+    weight_versions: list[WeightVersionSpan] | None = None
 
 
 class CompletionResponse(OpenAIBaseModel):
@@ -716,6 +728,8 @@ class CompletionResponseStreamChoice(OpenAIBaseModel):
     # prompt tokens is put into choice to align with CompletionResponseChoice
     prompt_token_ids: list[int] | None = None
     token_ids: list[int] | None = None
+    # Set on the final chunk only; covers the whole output.
+    weight_versions: list[WeightVersionSpan] | None = None
 
 
 class CompletionStreamResponse(OpenAIBaseModel):

vllm/entrypoints/openai/completion/serving.py

@@ -376,6 +376,7 @@ class OpenAIServingCompletion(GenerateBaseServing):
                             not delta_text
                             and not delta_token_ids
                             and not previous_num_tokens[i]
+                            and finish_reason is None
                         ):
                             # Chunked prefill case, don't return empty chunks
                             continue
@@ -421,6 +422,11 @@ class OpenAIServingCompletion(GenerateBaseServing):
                                     if request.return_token_ids
                                     else None
                                 ),
+                                weight_versions=(
+                                    output.weight_versions
+                                    if request.return_weight_versions
+                                    else None
+                                ),
                             )
                         ],
                     )
@@ -594,6 +600,11 @@ class OpenAIServingCompletion(GenerateBaseServing):
                         as_list(output.token_ids) if request.return_token_ids else None
                     ),
                     routed_experts=routed_experts_b64,
+                    weight_versions=(
+                        output.weight_versions
+                        if request.return_weight_versions
+                        else None
+                    ),
                 )
                 choices.append(choice_data)