From 463be19d83721ec5a63ee14905a65557d9a38d01 Mon Sep 17 00:00:00 2001 From: chengcuiping Date: Wed, 26 Aug 2026 21:13:45 +0800 Subject: [PATCH] [Core][V1] Reuse request token storage for trace replay Co-authored-by: OpenAI Codex Signed-off-by: chengcuiping --- tests/v1/worker/test_gpu_trace_replay.py | 233 ++++++++++++++++------ vllm/v1/worker/gpu/model_runner.py | 12 ++ vllm/v1/worker/gpu/sample/sampler.py | 12 ++ vllm/v1/worker/gpu/sample/trace_replay.py | 45 +++-- vllm/v1/worker/gpu/states.py | 4 + 5 files changed, 228 insertions(+), 78 deletions(-) diff --git a/tests/v1/worker/test_gpu_trace_replay.py b/tests/v1/worker/test_gpu_trace_replay.py index 9d418101fce4..708be526e89a 100644 --- a/tests/v1/worker/test_gpu_trace_replay.py +++ b/tests/v1/worker/test_gpu_trace_replay.py @@ -8,9 +8,6 @@ overwrites the sampled token in place before logprobs are computed. """ -from types import SimpleNamespace -from typing import cast - import pytest import torch @@ -38,28 +35,46 @@ def _i64(x) -> torch.Tensor: def _trace_state(max_num_reqs: int) -> TraceReplayState: - """Build a state whose req_states exposes the buffers apply_trace reads. - - total_len/prompt_len are written by the caller to position each request at - the desired replay step. - """ - req_states = cast( - RequestState, - SimpleNamespace( - max_num_reqs=max_num_reqs, - max_model_len=TEST_MAX_MODEL_LEN, - device=torch.device(DEVICE), - total_len=SimpleNamespace( - gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) - ), - prompt_len=SimpleNamespace( - gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE) - ), - ), + req_states = RequestState( + max_num_reqs=max_num_reqs, + max_model_len=TEST_MAX_MODEL_LEN, + max_num_batched_tokens=TEST_MAX_MODEL_LEN, + num_speculative_steps=0, + vocab_size=1024, + device=torch.device(DEVICE), ) return TraceReplayState(req_states) +def _admit( + state: TraceReplayState, + req_id: str, + prompt: list[int], + params: SamplingParams, + prefill: list[int] | None = None, + num_computed_tokens: int = 0, +) -> int: + prefill = prompt if prefill is None else prefill + req_states = state.req_states + suffix = state.get_token_suffix(len(prompt), len(prefill), params) + req_states.add_request( + req_id=req_id, + prompt_len=len(prompt), + all_token_ids=prefill, + num_computed_tokens=num_computed_tokens, + max_tokens=params.max_tokens, + future_token_ids=suffix, + ) + req_idx = req_states.req_id_to_index[req_id] + state.add_request(req_idx, params) + return req_idx + + +def _apply_admissions(state: TraceReplayState) -> None: + state.req_states.apply_staged_writes() + state.apply_staged_writes() + + def _set_lens(state: TraceReplayState, total_len, prompt_len) -> None: """Position each request at a replay step via total_len - prompt_len.""" state.req_states.total_len.gpu[: len(total_len)] = _i32(total_len) @@ -73,11 +88,10 @@ def test_replay_overwrites_sampled_at_each_step(): """The trace token for the current step replaces the sampled token.""" trace = [[100, 101, 102], [200, 201, 202]] trace_len = _i32([3, 3]) - trace_token_ids = torch.zeros( - 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE - ) + all_token_ids = torch.zeros(2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE) for i, t in enumerate(trace): - trace_token_ids[i, : len(t)] = _i32(t) + start = (5, 8)[i] + all_token_ids[i, start : start + len(t)] = _i32(t) prompt_len = _i32([5, 8]) idx_mapping = _i32([0, 1]) @@ -85,17 +99,15 @@ def test_replay_overwrites_sampled_at_each_step(): sampled = _i64([-7, -7]) # sentinel that must be overwritten total_len = _i32([5 + step, 8 + step]) apply_trace_tokens( - sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + sampled, idx_mapping, all_token_ids, trace_len, total_len, prompt_len ) assert sampled.tolist() == [trace[0][step], trace[1][step]] def test_past_end_of_trace_leaves_sampled_untouched(): """Once step >= trace_len, the sampler's own token is kept.""" - trace_token_ids = torch.zeros( - 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE - ) - trace_token_ids[0, :2] = _i32([100, 101]) + all_token_ids = torch.zeros(1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE) + all_token_ids[0, 4:6] = _i32([100, 101]) trace_len = _i32([2]) prompt_len = _i32([4]) idx_mapping = _i32([0]) @@ -103,16 +115,14 @@ def test_past_end_of_trace_leaves_sampled_untouched(): sampled = _i64([999]) total_len = _i32([4 + 2]) # step == 2 == trace_len -> out of range apply_trace_tokens( - sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + sampled, idx_mapping, all_token_ids, trace_len, total_len, prompt_len ) assert sampled.tolist() == [999] def test_non_trace_request_untouched(): """trace_len == 0 means the request never uses replay.""" - trace_token_ids = torch.zeros( - 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE - ) + all_token_ids = torch.zeros(1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE) trace_len = _i32([0]) prompt_len = _i32([3]) idx_mapping = _i32([0]) @@ -120,7 +130,7 @@ def test_non_trace_request_untouched(): sampled = _i64([42]) total_len = _i32([3]) apply_trace_tokens( - sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + sampled, idx_mapping, all_token_ids, trace_len, total_len, prompt_len ) assert sampled.tolist() == [42] @@ -128,10 +138,8 @@ def test_non_trace_request_untouched(): def test_idx_mapping_indirection_and_negative_skip(): """batch_idx -> req_state_idx indirection, and negative entries are skipped.""" # req_state 0: no trace; req_state 1: trace [500, 501]. - trace_token_ids = torch.zeros( - 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE - ) - trace_token_ids[1, :2] = _i32([500, 501]) + all_token_ids = torch.zeros(2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE) + all_token_ids[1, 6:8] = _i32([500, 501]) trace_len = _i32([0, 2]) prompt_len = _i32([10, 6]) @@ -141,7 +149,7 @@ def test_idx_mapping_indirection_and_negative_skip(): # Indexed by req_state_idx: state 0 at total_len=10, state 1 at step 1. total_len = _i32([10, 6 + 1]) apply_trace_tokens( - sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len + sampled, idx_mapping, all_token_ids, trace_len, total_len, prompt_len ) # batch 0 -> state 1 step 1 -> 501; batch 1 masked; batch 2 -> state 0 no trace. assert sampled.tolist() == [501, 2, 3] @@ -153,13 +161,22 @@ def test_idx_mapping_indirection_and_negative_skip(): def test_state_end_to_end(): """add_request -> apply_staged_writes -> apply_trace overwrites correctly.""" state = _trace_state(4) - state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22, 33])) - state.add_request(1, SamplingParams()) # no trace - state.apply_staged_writes() + prompt = [1, 2, 3, 4, 5, 6, 7] + trace_idx = _admit( + state, + "trace", + prompt, + SamplingParams(trace_decode_token_ids=[11, 22, 33]), + ) + normal_idx = _admit(state, "normal", prompt, SamplingParams()) + _apply_admissions(state) - idx_mapping = _i32([0, 1]) + assert not hasattr(state, "trace_token_ids") + assert state.req_states.all_token_ids.gpu[trace_idx, 7:10].tolist() == [11, 22, 33] + + idx_mapping = _i32([trace_idx, normal_idx]) sampled = _i64([-1, -1]) - _set_lens(state, [7 + 1, 7, 0, 0], [7, 7, 0, 0]) # req 0 at step 1 + state.req_states.total_len.gpu[trace_idx] = 7 + 1 state.apply_trace(sampled, idx_mapping) assert sampled.tolist() == [22, -1] @@ -167,29 +184,127 @@ def test_state_end_to_end(): def test_state_leaves_non_trace_batch_unchanged(): """Requests with trace_len == 0 remain unchanged after a trace was seen.""" state = _trace_state(4) - state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22])) - state.add_request(1, SamplingParams()) - state.apply_staged_writes() + prompt = [1, 2, 3, 4, 5, 6, 7] + _admit( + state, + "trace", + prompt, + SamplingParams(trace_decode_token_ids=[11, 22]), + ) + normal_idx = _admit(state, "normal", prompt, SamplingParams()) + _apply_admissions(state) - # Only the non-trace request (state 1) is in this batch. - idx_mapping = _i32([1]) + idx_mapping = _i32([normal_idx]) sampled = _i64([555]) - _set_lens(state, [0, 7, 0, 0], [7, 7, 0, 0]) state.apply_trace(sampled, idx_mapping) assert sampled.tolist() == [555] def test_slot_reuse_clears_trace(): """Reusing a slot for a non-trace request must not replay stale tokens.""" - state = _trace_state(2) - state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22])) - state.apply_staged_writes() - # Slot 0 reused by a request without a trace. - state.add_request(0, SamplingParams()) - state.apply_staged_writes() + state = _trace_state(1) + trace_idx = _admit( + state, + "trace", + [1, 2, 3], + SamplingParams(trace_decode_token_ids=[11, 22]), + ) + _apply_admissions(state) + assert state.req_states.remove_request("trace") == trace_idx + normal_idx = _admit(state, "normal", [7, 8, 9], SamplingParams()) + _apply_admissions(state) + assert normal_idx == trace_idx - idx_mapping = _i32([0]) + idx_mapping = _i32([normal_idx]) sampled = _i64([888]) - _set_lens(state, [3, 0], [3, 0]) state.apply_trace(sampled, idx_mapping) assert sampled.tolist() == [888] + + +def test_resume_stages_only_unconsumed_trace_suffix(): + """A resumed request derives its trace offset from its prefill history.""" + state = _trace_state(2) + prompt = [1, 2, 3] + trace = [11, 22, 33] + req_idx = _admit( + state, + "resumed", + prompt, + SamplingParams(trace_decode_token_ids=trace), + prefill=prompt + trace[:1], + ) + _apply_admissions(state) + + assert state.req_states.all_token_ids.gpu[req_idx, :6].tolist() == prompt + trace + sampled = _i64([-1]) + state.apply_trace(sampled, _i32([req_idx])) + assert sampled.tolist() == [22] + + +def test_full_admission_batch_uses_one_staged_write_per_request(): + """Trace suffixes must not overflow staged-write metadata at full capacity.""" + state = _trace_state(2) + params = SamplingParams(trace_decode_token_ids=[11, 22]) + first_idx = _admit(state, "first", [1, 2, 3], params) + second_idx = _admit(state, "second", [4, 5, 6], params) + + assert len(state.req_states.all_token_ids._staged_write_indices) == 2 + _apply_admissions(state) + assert state.req_states.all_token_ids.gpu[first_idx, 3:5].tolist() == [11, 22] + assert state.req_states.all_token_ids.gpu[second_idx, 3:5].tolist() == [11, 22] + + +def test_partial_prefill_admission_preserves_trace_start(): + state = _trace_state(1) + prompt = [1, 2, 3, 4] + req_idx = _admit( + state, + "partial-prefill", + prompt, + SamplingParams(trace_decode_token_ids=[11, 22]), + num_computed_tokens=2, + ) + _apply_admissions(state) + + assert state.req_states.num_computed_tokens.gpu[req_idx].item() == 2 + assert state.req_states.all_token_ids.gpu[req_idx, :6].tolist() == [ + *prompt, + 11, + 22, + ] + + +def test_exhausted_trace_adds_no_future_suffix(): + state = _trace_state(1) + prompt = [1, 2, 3] + trace = [11, 22] + req_idx = _admit( + state, + "exhausted", + prompt, + SamplingParams(trace_decode_token_ids=trace), + prefill=prompt + trace, + ) + + assert len(state.req_states.all_token_ids._staged_write_indices) == 1 + assert state.req_states.all_token_ids._staged_write_contents == prompt + trace + _apply_admissions(state) + + sampled = _i64([999]) + state.apply_trace(sampled, _i32([req_idx])) + assert sampled.tolist() == [999] + + +def test_empty_future_suffix_preserves_empty_write_behavior(): + req_states = _trace_state(1).req_states + req_states.add_request( + req_id="empty", + prompt_len=0, + all_token_ids=[], + num_computed_tokens=0, + max_tokens=1, + future_token_ids=[], + ) + + assert req_states.all_token_ids._staged_write_indices == [] + assert req_states.all_token_ids._staged_write_contents == [] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3bdccebfe2ce..db9cd119d311 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -987,12 +987,24 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: prompt_len = new_req_data.prompt_len sampling_params = new_req_data.sampling_params + trace_token_suffix = None + if ( + self.is_last_pp_rank + and sampling_params is not None + and self.sampler is not None + ): + trace_token_suffix = self.sampler.get_trace_token_suffix( + prompt_len, + len(new_req_data.prefill_token_ids), + sampling_params, + ) self.req_states.add_request( req_id=req_id, prompt_len=prompt_len, all_token_ids=new_req_data.prefill_token_ids, num_computed_tokens=new_req_data.num_computed_tokens, max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type] + future_token_ids=trace_token_suffix, ) req_index = self.req_states.req_id_to_index[req_id] if self.adaptive_verification is not None: diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index d3205a3dae83..06b75b3158aa 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -65,6 +65,18 @@ def __init__( not return_sampling_mask and flashinfer_sampler_supported() ) + def get_trace_token_suffix( + self, + prompt_len: int, + prefill_len: int, + sampling_params: SamplingParams, + ) -> list[int] | None: + if self.trace_replay_state is None: + return None + return self.trace_replay_state.get_token_suffix( + prompt_len, prefill_len, sampling_params + ) + def add_request( self, req_idx: int, prompt_len: int, sampling_params: SamplingParams ) -> None: diff --git a/vllm/v1/worker/gpu/sample/trace_replay.py b/vllm/v1/worker/gpu/sample/trace_replay.py index 9b58d4cf2923..15c3c125b08d 100644 --- a/vllm/v1/worker/gpu/sample/trace_replay.py +++ b/vllm/v1/worker/gpu/sample/trace_replay.py @@ -4,7 +4,7 @@ from vllm.sampling_params import SamplingParams from vllm.triton_utils import tl, triton -from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.states import RequestState @@ -23,25 +23,33 @@ def __init__(self, req_states: RequestState): self.max_num_reqs = req_states.max_num_reqs self.device = req_states.device self.req_states = req_states - self.trace_token_ids = StagedWriteTensor( - (self.max_num_reqs, req_states.max_model_len), - dtype=torch.int32, - device=self.device, - uva_instead_of_gpu=True, - ) self.trace_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32) def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: trace = sampling_params.trace_decode_token_ids if trace is not None: self.trace_len.np[req_idx] = len(trace) - self.trace_token_ids.stage_write(req_idx, 0, trace) else: self.trace_len.np[req_idx] = 0 + def get_token_suffix( + self, + prompt_len: int, + prefill_len: int, + sampling_params: SamplingParams, + ) -> list[int]: + """Return unconsumed trace tokens to append to all_token_ids.""" + trace = sampling_params.trace_decode_token_ids + if trace is None: + return [] + trace_offset = prefill_len - prompt_len + assert 0 <= trace_offset <= len(trace) + remaining_trace = trace[trace_offset:] + assert prefill_len + len(remaining_trace) <= self.req_states.max_model_len + return remaining_trace + def apply_staged_writes(self) -> None: self.trace_len.copy_to_uva() - self.trace_token_ids.apply_write() def apply_trace( self, @@ -51,7 +59,7 @@ def apply_trace( apply_trace_tokens( sampled, idx_mapping, - self.trace_token_ids.gpu, + self.req_states.all_token_ids.gpu, self.trace_len.gpu, self.req_states.total_len.gpu, self.req_states.prompt_len.gpu, @@ -62,8 +70,8 @@ def apply_trace( def _trace_replay_kernel( sampled_ptr, # [num_reqs], int64, mutated in place idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx - trace_token_ids_ptr, # [max_num_reqs, max_model_len], int32 - trace_token_ids_stride, + all_token_ids_ptr, # [max_num_reqs, max_model_len], int32 + all_token_ids_stride, trace_len_ptr, # [max_num_reqs], int32 total_len_ptr, # [max_num_reqs], int32 prompt_len_ptr, # [max_num_reqs], int32 @@ -80,14 +88,13 @@ def _trace_replay_kernel( # The token being sampled now is output token number # (total_len - prompt_len): total_len reflects tokens committed through the # previous step (post_update runs after sampling). - step = tl.load(total_len_ptr + req_state_idx) - tl.load( - prompt_len_ptr + req_state_idx - ) + prompt_len = tl.load(prompt_len_ptr + req_state_idx) + step = tl.load(total_len_ptr + req_state_idx) - prompt_len if step < 0 or step >= trace_len: return token_id = tl.load( - trace_token_ids_ptr + req_state_idx * trace_token_ids_stride + step + all_token_ids_ptr + req_state_idx * all_token_ids_stride + prompt_len + step ) tl.store(sampled_ptr + batch_idx, token_id.to(tl.int64)) @@ -95,7 +102,7 @@ def _trace_replay_kernel( def apply_trace_tokens( sampled: torch.Tensor, idx_mapping: torch.Tensor, - trace_token_ids: torch.Tensor, + all_token_ids: torch.Tensor, trace_len: torch.Tensor, total_len: torch.Tensor, prompt_len: torch.Tensor, @@ -105,8 +112,8 @@ def apply_trace_tokens( _trace_replay_kernel[(num_reqs,)]( sampled, idx_mapping, - trace_token_ids, - trace_token_ids.stride(0), + all_token_ids, + all_token_ids.stride(0), trace_len, total_len, prompt_len, diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 025c8300fcd2..9bfd00e6f96c 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -95,6 +95,7 @@ def add_request( all_token_ids: list[int], num_computed_tokens: int, max_tokens: int, + future_token_ids: list[int] | None = None, ) -> None: assert len(self.free_indices) > 0, "No free indices" req_idx = self.free_indices.pop() @@ -109,6 +110,9 @@ def add_request( ) self.prefill_len.np[req_idx] = prefill_len self.total_len.stage_write_elem(req_idx, prefill_len) + if future_token_ids: + assert prefill_len + len(future_token_ids) <= self.max_model_len + all_token_ids = all_token_ids + future_token_ids self.all_token_ids.stage_write(req_idx, 0, all_token_ids) self.num_computed_prefill_tokens[req_idx] = num_computed_tokens self.num_computed_tokens_np[req_idx] = num_computed_tokens