From e724154842858712d1108dc3329e0d45aaa807a9 Mon Sep 17 00:00:00 2001 From: Shuolei Wang Date: Thu, 1 Oct 2026 17:38:25 +0800 Subject: [PATCH 1/2] [Core][Frontend] Return token-level weight-version spans on final outputs EngineCore stamps its committed weight-version label on every scheduled step, and the step's EngineCoreOutputs carries it. OutputProcessor records a span start whenever a request's tokens arrive under a new label, and the final output of every sequence returns `weight_versions: [{version, start, end}]` on CompletionOutput and on /inference/v1/generate choices. Signed-off-by: Shuolei Wang --- .../engine-core-client/src/protocol/output.rs | 4 ++ .../src/tests/python_compat.py | 3 + tests/v1/engine/test_output_processor.py | 55 +++++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 6 ++ .../scale_out/token_in_token_out/serving.py | 2 + vllm/outputs.py | 15 +++++ vllm/v1/core/sched/output.py | 4 ++ vllm/v1/core/sched/scheduler.py | 4 +- vllm/v1/engine/__init__.py | 3 + vllm/v1/engine/async_llm.py | 5 +- vllm/v1/engine/core.py | 2 + vllm/v1/engine/llm_engine.py | 1 + vllm/v1/engine/output_processor.py | 28 ++++++++++ 13 files changed, 130 insertions(+), 2 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 9339c117a889..71ac276c0844 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -228,6 +228,10 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, + /// Weight-version label of the step that produced these outputs. Decoded to + /// match the Python schema; the Rust frontend does not surface it yet. + #[serde(default)] + weight_version: Option, } /// Data-parallel control notifications multiplexed through `EngineCoreOutputs`. diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index bd8d40dc32f2..c0ca97dbe638 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/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. diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 44c3ee20b5e4..af7d355a6dd0 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/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,55 @@ 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", + [ + # Same-label steps share a span, a step can return several accepted + # tokens, and labels are opaque: A -> B -> A gives 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 + + # The abort output built by the frontend carries the spans. Unread deltas + # are merged into it, and the offsets still cover the whole output. + 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] diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 84a6a2c62388..46a902493ced 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/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,9 @@ 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 + # Weight versions that sampled the output tokens, as contiguous spans over + # output-token indexes, in generation order. + weight_versions: list[WeightVersionSpan] | None = None @field_validator("token_ids") @classmethod @@ -443,6 +447,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): diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index b6d7bf19ba93..48e3466870f8 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/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: diff --git a/vllm/outputs.py b/vllm/outputs.py index 8e69ffa91d91..1dc99b55adf4 100644 --- a/vllm/outputs.py +++ b/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,9 @@ 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: The weight versions that sampled the output tokens, as + contiguous spans over output-token indexes, in generation order. + Set on the final output only. """ @@ -76,6 +88,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 +236,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: diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index a5af5a3d5755..43b85a3b16ca 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -311,6 +311,10 @@ class SchedulerOutput: # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 + # Label of the weights this step runs with, stamped by EngineCore right + # after scheduling. Tokens sampled in this step are attributed to it. + weight_version: str | None = None + @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 0758c6122565..11f8c520b36f 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2386,7 +2386,9 @@ def update_from_output( # Create EngineCoreOutputs for all clients that have requests with # outputs in this step. engine_core_outputs = { - client_index: EngineCoreOutputs(outputs=outs) + client_index: EngineCoreOutputs( + outputs=outs, weight_version=scheduler_output.weight_version + ) for client_index, outs in outputs.items() } diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index effe8d2416e4..0418e27eb937 100644 --- a/vllm/v1/engine/__init__.py +++ b/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 + # Label of the weights that ran the step these outputs came from. + weight_version: str | None = None + def __post_init__(self): if self.timestamp == 0.0: self.timestamp = time.monotonic() diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7e50aa045123..02bebda8cb8f 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -852,7 +852,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, + outputs.weight_version, ) # NOTE: RequestOutputs are pushed to their queues. assert not processed_outputs.request_outputs diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b13764abec3a..4186948072c0 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -655,6 +655,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: if not self.scheduler.has_requests(): return {}, False scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) + scheduler_output.weight_version = self._weight_version future = self.model_executor.execute_model(scheduler_output, non_block=True) grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) with ( @@ -712,6 +713,7 @@ def step_with_batch_queue( deferred_scheduler_output = None if self.scheduler.has_requests(): scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) + scheduler_output.weight_version = self._weight_version with self.log_error_detail(scheduler_output): exec_future = self.model_executor.execute_model( scheduler_output, non_block=True diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 2dd70fc17917..330c114a86dd 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -335,6 +335,7 @@ def step(self) -> list[RequestOutput | PoolingRequestOutput]: outputs.outputs, engine_core_timestamp=outputs.timestamp, iteration_stats=iteration_stats, + weight_version=outputs.weight_version, ) mm_processor_cache = self.renderer.mm_processor_cache diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index abb45f7239be..73e8681aa5ae 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -19,6 +19,7 @@ RequestError, RequestOutput, SamplingMask, + WeightVersionSpan, ) from vllm.sampling_params import RequestOutputKind from vllm.tokenizers import TokenizerLike @@ -190,6 +191,8 @@ def __init__( # Routed experts accumulation (prompt + sample chunks) self.routed_experts_chunks: list[np.ndarray] = [] self.sampling_mask_chunks: list[SamplingMaskLists] = [] + # (weight version, first output-token index) for each span. + self.weight_version_starts: list[tuple[str, int]] = [] # Stream Interval self.stream_interval = stream_interval @@ -460,6 +463,18 @@ def _new_completion_output( if finished and self.routed_experts_chunks: routed_experts = np.concatenate(self.routed_experts_chunks, axis=0) + # Final outputs carry the spans; a zero-token output carries []. + weight_versions = None + if finished: + num_output_tokens = self.detokenizer.num_output_tokens() + starts = self.weight_version_starts + if starts or num_output_tokens == 0: + ends = [start for _, start in starts[1:]] + [num_output_tokens] + weight_versions = [ + WeightVersionSpan(version, start, end) + for (version, start), end in zip(starts, ends) + ] + return CompletionOutput( index=self.request_index, text=text, @@ -471,6 +486,7 @@ def _new_completion_output( finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, spec_decode_metrics=self.spec_decode_metrics if finished else None, + weight_versions=weight_versions, ) def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: @@ -659,6 +675,7 @@ def process_outputs( engine_core_outputs: list[EngineCoreOutput], engine_core_timestamp: float | None = None, iteration_stats: IterationStats | None = None, + weight_version: str | None = None, ) -> OutputProcessorOutput: """Process the EngineCoreOutputs: 1) Compute stats for logging @@ -736,6 +753,17 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None + # A new span starts where this step's label differs from the + # label of the request's previous tokens. + starts = req_state.weight_version_starts + if ( + new_token_ids + and weight_version is not None + and (not starts or starts[-1][0] != weight_version) + ): + starts.append( + (weight_version, req_state.detokenizer.num_output_tokens()) + ) # 2) Detokenize the token ids into text and perform stop checks. num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update( From 006ebac240f34ddcd36383832b28c25840d2ea80 Mon Sep 17 00:00:00 2001 From: Shuolei Wang Date: Sat, 10 Oct 2026 01:49:31 +0800 Subject: [PATCH 2/2] Stamp weight version in EngineCore, add docs - EngineCore stamps each step's outputs; the scheduler is unchanged. - Span tracking moves into two small RequestState helpers. - Document weight_versions and use it in the async RL example. Signed-off-by: Shuolei Wang Co-authored-by: Codex Signed-off-by: Shuolei Wang <948904026@qq.com> --- .../online_serving/token_in_token_out.md | 15 +++++++ docs/training/async_rl.md | 20 ++++++++- examples/rl/rlhf_async_new_apis.py | 43 ++++++++----------- .../engine-core-client/src/protocol/output.rs | 3 +- tests/v1/engine/test_output_processor.py | 6 +-- .../scale_out/token_in_token_out/protocol.py | 3 +- vllm/outputs.py | 3 +- vllm/v1/core/sched/output.py | 4 -- vllm/v1/core/sched/scheduler.py | 4 +- vllm/v1/engine/__init__.py | 2 +- vllm/v1/engine/async_llm.py | 2 +- vllm/v1/engine/core.py | 9 +++- vllm/v1/engine/output_processor.py | 43 ++++++++----------- 13 files changed, 85 insertions(+), 72 deletions(-) diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md index e22d4906155b..b3e15ae61692 100644 --- a/docs/serving/online_serving/token_in_token_out.md +++ b/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`. diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index 9e75a24eaa1a..b5ae067e3cb1 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -43,10 +43,11 @@ 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 @@ -54,12 +55,27 @@ 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. diff --git a/examples/rl/rlhf_async_new_apis.py b/examples/rl/rlhf_async_new_apis.py index 081023156edf..3449d0f17f20 100644 --- a/examples/rl/rlhf_async_new_apis.py +++ b/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 diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 71ac276c0844..8945f8b11020 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -228,8 +228,7 @@ struct WireEngineCoreOutputs { /// wave needs to start in other engines. #[serde(default)] start_wave: Option, - /// Weight-version label of the step that produced these outputs. Decoded to - /// match the Python schema; the Rust frontend does not surface it yet. + /// Decoded for Python wire compatibility; not exposed by the Rust frontend. #[serde(default)] weight_version: Option, } diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index af7d355a6dd0..8eb85765993b 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -1736,8 +1736,7 @@ def test_async_pooling_cache_miss_does_not_fail_output_processor(): @pytest.mark.parametrize( "steps,expected", [ - # Same-label steps share a span, a step can return several accepted - # tokens, and labels are opaque: A -> B -> A gives three spans. + # 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)], @@ -1774,8 +1773,7 @@ def test_weight_version_spans(output_kind: RequestOutputKind, steps, expected): assert intermediate is not None assert intermediate.outputs[0].weight_versions is None - # The abort output built by the frontend carries the spans. Unread deltas - # are merged into it, and the offsets still cover the whole output. + # 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 diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 46a902493ced..bccc8d8eff0b 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -412,8 +412,7 @@ 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 - # Weight versions that sampled the output tokens, as contiguous spans over - # output-token indexes, in generation order. + # Contiguous weight version spans over output-token indexes. weight_versions: list[WeightVersionSpan] | None = None @field_validator("token_ids") diff --git a/vllm/outputs.py b/vllm/outputs.py index 1dc99b55adf4..145b7d510f71 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -71,8 +71,7 @@ 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: The weight versions that sampled the output tokens, as - contiguous spans over output-token indexes, in generation order. + weight_versions: Contiguous weight version spans over output-token indexes. Set on the final output only. """ diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 43b85a3b16ca..a5af5a3d5755 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -311,10 +311,6 @@ class SchedulerOutput: # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 - # Label of the weights this step runs with, stamped by EngineCore right - # after scheduling. Tokens sampled in this step are attributed to it. - weight_version: str | None = None - @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 11f8c520b36f..0758c6122565 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2386,9 +2386,7 @@ def update_from_output( # Create EngineCoreOutputs for all clients that have requests with # outputs in this step. engine_core_outputs = { - client_index: EngineCoreOutputs( - outputs=outs, weight_version=scheduler_output.weight_version - ) + client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items() } diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 0418e27eb937..1b14b08246aa 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -282,7 +282,7 @@ class EngineCoreOutputs( # "old" wave, so the next wave needs to be started in other engines. start_wave: int | None = None - # Label of the weights that ran the step these outputs came from. + # Weight version used for this step. weight_version: str | None = None def __post_init__(self): diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 02bebda8cb8f..dd73e2de587a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -855,7 +855,7 @@ async def output_handler(): outputs_slice, outputs.timestamp, iteration_stats, - outputs.weight_version, + weight_version=outputs.weight_version, ) # NOTE: RequestOutputs are pushed to their queues. assert not processed_outputs.request_outputs diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 4186948072c0..bbca64a58ae9 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -639,6 +639,11 @@ def _attach_iteration_details( else: eco.scheduler_stats.iteration_details = iteration_details + def _attach_weight_version(self, outputs: dict[int, EngineCoreOutputs]) -> None: + # Weight versions change only while generation is paused. + for eco in outputs.values(): + eco.weight_version = self._weight_version + def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). Overridden by the DP engine core; never throttles otherwise.""" @@ -655,7 +660,6 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: if not self.scheduler.has_requests(): return {}, False scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) - scheduler_output.weight_version = self._weight_version future = self.model_executor.execute_model(scheduler_output, non_block=True) grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) with ( @@ -673,6 +677,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 @@ -713,7 +718,6 @@ def step_with_batch_queue( deferred_scheduler_output = None if self.scheduler.has_requests(): scheduler_output = self.scheduler.schedule(self._should_throttle_prefills()) - scheduler_output.weight_version = self._weight_version with self.log_error_detail(scheduler_output): exec_future = self.model_executor.execute_model( scheduler_output, non_block=True @@ -775,6 +779,7 @@ def step_with_batch_queue( scheduler_output, model_output ) self._attach_iteration_details(engine_core_outputs, iteration_details) + self._attach_weight_version(engine_core_outputs) # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 73e8681aa5ae..263a28fb3901 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -463,18 +463,6 @@ def _new_completion_output( if finished and self.routed_experts_chunks: routed_experts = np.concatenate(self.routed_experts_chunks, axis=0) - # Final outputs carry the spans; a zero-token output carries []. - weight_versions = None - if finished: - num_output_tokens = self.detokenizer.num_output_tokens() - starts = self.weight_version_starts - if starts or num_output_tokens == 0: - ends = [start for _, start in starts[1:]] + [num_output_tokens] - weight_versions = [ - WeightVersionSpan(version, start, end) - for (version, start), end in zip(starts, ends) - ] - return CompletionOutput( index=self.request_index, text=text, @@ -486,9 +474,25 @@ def _new_completion_output( finish_reason=str(finish_reason) if finished else None, stop_reason=stop_reason if finished else None, spec_decode_metrics=self.spec_decode_metrics if finished else None, - weight_versions=weight_versions, + weight_versions=self._weight_version_spans() if finished else None, ) + def track_weight_version(self, weight_version: str) -> None: + """Start a span when the weight version changes.""" + assert self.detokenizer is not None + starts = self.weight_version_starts + if not starts or starts[-1][0] != weight_version: + starts.append((weight_version, self.detokenizer.num_output_tokens())) + + def _weight_version_spans(self) -> list[WeightVersionSpan] | None: + assert self.detokenizer is not None + num_tokens = self.detokenizer.num_output_tokens() + starts = self.weight_version_starts + if not starts and num_tokens: + return None # No label was reported for these tokens. + ends = [start for _, start in starts[1:]] + [num_tokens] + return [WeightVersionSpan(v, s, e) for (v, s), e in zip(starts, ends)] + def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput: return PoolingOutput(data=pooling_output) @@ -753,17 +757,8 @@ def process_outputs( if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None - # A new span starts where this step's label differs from the - # label of the request's previous tokens. - starts = req_state.weight_version_starts - if ( - new_token_ids - and weight_version is not None - and (not starts or starts[-1][0] != weight_version) - ): - starts.append( - (weight_version, req_state.detokenizer.num_output_tokens()) - ) + if new_token_ids and weight_version is not None: + req_state.track_weight_version(weight_version) # 2) Detokenize the token ids into text and perform stop checks. num_prev_tokens = req_state.detokenizer.num_output_tokens() stop_string = req_state.detokenizer.update(