[Bugfix][KV Offload] Reset a resumed streaming session's block table in OffloadingConnector and SimpleCPUOffloadConnector

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -1190,6 +1190,24 @@ def take_raw_events():
     assert events[0].token_ids == token_ids
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_streaming_session_stores_new_chunks(request_runner, async_scheduling: bool):
+    """A streaming session resumed for its next input chunk arrives with its
+    full block table, so its new chunks must be stored from their own blocks."""
+    runner = request_runner(
+        block_size=4, num_gpu_blocks=100, async_scheduling=async_scheduling
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.new_request(token_ids=list(range(1, 9)), req_id="s", resumable=True)
+    runner.run(decoded_tokens=[100, EOS_TOKEN_ID], expected_stored=(0, 1))
+
+    runner.new_request(token_ids=list(range(200, 212)), req_id="s", resumable=True)
+    runner.run(decoded_tokens=[300, EOS_TOKEN_ID], expected_stored=(2, 3, 4))
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_offloading_connector(request_runner, async_scheduling: bool):
     block_size = 4

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -369,8 +369,12 @@ def new_request(
         token_ids: list[int],
         kv_transfer_params: dict | None = None,
         skip_reading_prefix_cache: bool = False,
+        req_id: str | None = None,
+        resumable: bool = False,
     ):
-        self.req_id += 1
+        if req_id is None:
+            self.req_id += 1
+            req_id = str(self.req_id)
 
         sampling_params = SamplingParams(
             max_tokens=1000,
@@ -379,11 +383,12 @@ def new_request(
         sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
 
         req = Request(
-            request_id=str(self.req_id),
+            request_id=req_id,
             prompt_token_ids=token_ids,
             sampling_params=sampling_params,
             pooling_params=None,
             block_hasher=self._block_hasher,
+            resumable=resumable,
         )
         if kv_transfer_params is not None:
             req.kv_transfer_params = kv_transfer_params

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -296,12 +296,14 @@ def yield_req_data(
     scheduler_output,
 ) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]:
     """Yields:
-    (req_id, new_block_id_groups, preempted)
+    (req_id, block_id_groups, is_full_table)
 
+    ``is_full_table`` is True when ``block_id_groups`` is the request's whole
+    block table (new and resumed requests) rather than this step's additions.
     """
     # new requests
     for req_data in scheduler_output.scheduled_new_reqs:
-        yield req_data.req_id, req_data.block_ids, False
+        yield req_data.req_id, req_data.block_ids, True
 
     # cached requests
     cached_reqs = scheduler_output.scheduled_cached_reqs

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -1285,11 +1285,13 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
         # Used to detect sliding window blocks that got re-allocated.
         new_block_ids_end: dict[str, tuple[int, ...]] = {}
 
-        for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):
+        for req_id, new_block_id_groups, is_full_table in yield_req_data(
+            scheduler_output
+        ):
             req_status = self._req_status[req_id]
             req_status.update_offload_keys()
 
-            if preempted:
+            if is_full_table:
                 for group_state in req_status.group_states:
                     group_state.block_ids.clear()