diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index b5ae067e3c..6442595999 100644 --- a/docs/training/async_rl.md +++ b/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. diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 747b84a938..85404af06c 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/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 diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 17934026ef..0217d8a733 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/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, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 79dd9605e6..071f73f043 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/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, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9ddc20ed30..7e1782e241 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/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) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 4da67e607a..9e8ffb496b 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/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): diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 0306909237..944da27530 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/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)