Fix rollback of invalid output tokens and generator state

Signed-off-by: David Ben-David <davidb@pliops.com>

Author: David Ben-David

tests/v1/worker/test_gpu_model_runner.py

@@ -250,6 +250,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
         new_token_ids=[[]],
         new_block_ids=([[0]], ),
         num_computed_tokens=[0],
+        num_output_tokens=[0],
     )
 
     scheduler_output = SchedulerOutput(

vllm/v1/core/sched/output.py

@@ -101,6 +101,7 @@ class CachedRequestData:
     new_token_ids: list[list[int]]
     new_block_ids: list[Optional[tuple[list[int], ...]]]
     num_computed_tokens: list[int]
+    num_output_tokens: list[int]
 
     @property
     def num_reqs(self) -> int:
@@ -114,6 +115,7 @@ def make_empty(cls) -> CachedRequestData:
             new_token_ids=[],
             new_block_ids=[],
             num_computed_tokens=[],
+            num_output_tokens=[],
         )
 
 

vllm/v1/core/sched/scheduler.py

@@ -672,6 +672,7 @@ def _make_cached_request_data(
         new_token_ids: list[list[int]] = []
         new_block_ids: list[Optional[tuple[list[int], ...]]] = []
         num_computed_tokens: list[int] = []
+        num_output_tokens: list[int] = []
 
         use_connector = self.connector is not None
         for req in itertools.chain(running_reqs, resumed_reqs):
@@ -696,6 +697,7 @@ def _make_cached_request_data(
             new_block_ids.append(
                 req_to_new_blocks[req_id].get_block_ids(allow_none=True))
             num_computed_tokens.append(req.num_computed_tokens)
+            num_output_tokens.append(len(req.output_token_ids))
         # Because resumed_reqs is usually empty, it is more efficient to do
         # in-place appending so that we don't need to allocate a new list.
         resumed_from_preemption = [False] * len(running_reqs)
@@ -707,6 +709,7 @@ def _make_cached_request_data(
             new_token_ids=new_token_ids,
             new_block_ids=new_block_ids,
             num_computed_tokens=num_computed_tokens,
+            num_output_tokens=num_output_tokens,
         )
 
     def _try_schedule_encoder_inputs(

vllm/v1/worker/gpu_input_batch.py

@@ -48,11 +48,6 @@ class CachedRequestState:
     def __post_init__(self):
         self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
             self.prompt_token_ids, self.prompt_embeds)
-        # 'last_generator_offset' and 'len_last_output_token_ids' are used to
-        # allow safe rollback in case a sampled token turns out to be invalid
-        # (e.g., due to KV load errors).
-        self.last_generator_offset = 0 if self.generator else None
-        self.len_last_output_token_ids = len(self.output_token_ids)
 
     @property
     def num_tokens(self) -> int:
@@ -242,7 +237,6 @@ def __init__(
         # NOTE(woosuk): The indices of the requests that do not have their own
         # generator should not be included in the dictionary.
         self.generators: dict[int, torch.Generator] = {}
-        self.generators_last_offset: dict[int, int] = {}
 
         self.num_logprobs: dict[str, int] = {}
         # NOTE(rob): num_prompt_logprobs only includes reqs
@@ -393,9 +387,6 @@ def add_request(
             # do not have their own generator.
             if request.generator is not None:
                 self.generators[req_index] = request.generator
-                assert (request.last_generator_offset is not None)
-                self.generators_last_offset[
-                    req_index] = request.last_generator_offset
 
             if sampling_params.logprobs is not None:
                 self.num_logprobs[req_id] = (self.vocab_size

vllm/v1/worker/gpu_model_runner.py

@@ -634,25 +634,9 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
             num_computed_tokens = req_data.num_computed_tokens[i]
             new_block_ids = req_data.new_block_ids[i]
             resumed_from_preemption = req_data.resumed_from_preemption[i]
+            num_output_tokens = req_data.num_output_tokens[i]
 
             # Update the cached states.
-            if num_computed_tokens <= req_state.num_computed_tokens:
-                # The request was rescheduled after a KV load failure. Clear
-                # the last sampled tokens and rewind the generator state
-                len_output_token_ids = len(req_state.output_token_ids)
-                del req_state.output_token_ids[req_state.
-                                               len_last_output_token_ids:]
-                if req_state.generator:
-                    req_state.generator.set_offset(
-                        req_state.last_generator_offset)
-                req_index = self.input_batch.req_id_to_index.get(req_id)
-                if req_index is not None:
-                    len_last_sampled = (len_output_token_ids -
-                                        req_state.len_last_output_token_ids)
-                    end_idx = self.input_batch.num_tokens_no_spec[
-                        req_index] - len_last_sampled
-                    self.input_batch.num_tokens[req_index] = end_idx
-                    self.input_batch.num_tokens_no_spec[req_index] = end_idx
 
             req_state.num_computed_tokens = num_computed_tokens
 
@@ -671,12 +655,21 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
                 elif num_new_tokens > 0:
                     req_state.output_token_ids.extend(
                         new_token_ids[-num_new_tokens:])
+            elif num_output_tokens < len(req_state.output_token_ids):
+                # Some output tokens were discarded due to a sync-KV-load
+                # failure. Align the cached state.
+                del req_state.output_token_ids[num_output_tokens:]
 
-            req_state.len_last_output_token_ids = len(
-                req_state.output_token_ids)
-            if req_state.generator:
-                req_state.last_generator_offset = (
-                    req_state.generator.get_offset())
+                req_index = self.input_batch.req_id_to_index.get(req_id)
+                if req_index is not None:
+                    old_end_idx = self.input_batch.num_tokens_no_spec[
+                        req_index]
+                    end_idx = self.input_batch.num_prompt_tokens[
+                        req_index] + num_output_tokens
+                    self.input_batch.num_tokens[req_index] = end_idx
+                    self.input_batch.num_tokens_no_spec[req_index] = end_idx
+                    self.input_batch.is_token_ids[req_index,
+                                                  end_idx:old_end_idx] = False
 
             # Update the block IDs.
             if not resumed_from_preemption:
@@ -699,11 +692,6 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
                 reqs_to_add.append(req_state)
                 continue
 
-            if req_state.generator:
-                assert (req_state.last_generator_offset is not None)
-                self.input_batch.generators_last_offset[
-                    req_index] = req_state.last_generator_offset
-
             # Update the persistent batch.
             self.input_batch.num_computed_tokens_cpu[req_index] = (
                 num_computed_tokens)
@@ -2185,8 +2173,7 @@ def _bookkeeping_sync(
         for i in discard_sampled_tokens_req_indices:
             gen = self.input_batch.generators.get(int(i))
             if gen is not None:
-                offset = self.input_batch.generators_last_offset.get(int(i))
-                gen.set_offset(offset)
+                gen.set_offset(gen.get_offset() - 4)
 
         # Copy some objects so they don't get modified after returning.
         # This is important when using async scheduling.