[Bugfix][KV Offload] Preserve stored event prefix order

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

@@ -465,7 +465,7 @@ def test_lookup_promotion_factor_gt_1_store_and_remove():
     assert not tracker._pending_event_metadata
 
 
-def test_take_events_factor_gt_1_store_is_order_independent():
+def test_take_events_factor_gt_1_store_uses_prefix_order():
     blocks_per_chunk = 3
     tracker = _tracker()
     group_config = _group_config(blocks_per_chunk=blocks_per_chunk)
@@ -479,14 +479,35 @@ def test_take_events_factor_gt_1_store_is_order_independent():
     events = list(tracker.take_events([_stored_event([keys[1], unknown_key, keys[0]])]))
 
     assert len(events) == 3
-    chunk1, placeholder, chunk0 = events
-    assert [len(event.block_hashes) for event in events] == [3, 1, 3]
+    chunk0, chunk1, placeholder = events
+    assert [len(event.block_hashes) for event in events] == [3, 3, 1]
     assert placeholder.block_size == 0
     assert placeholder.token_ids == []
     assert chunk0.parent_block_hash is None
     assert chunk1.parent_block_hash == chunk0.block_hashes[-1]
 
 
+def test_take_events_sorts_store_keys_by_group_and_chunk():
+    tracker = _tracker()
+    req = _request(block_hashes=[_hash(0), _hash(1)], token_count=8)
+    group0 = _group_config(group_idx=0)
+    group1 = _group_config(group_idx=1)
+    keys0 = _record_chunks(tracker, req, group0, num_chunks=2)
+    keys1 = _record_chunks(tracker, req, group1, num_chunks=2)
+
+    events = list(
+        tracker.take_events([_stored_event([keys1[1], keys0[1], keys1[0], keys0[0]])])
+    )
+
+    assert [event.group_idx for event in events] == [0, 0, 1, 1]
+    assert [event.token_ids for event in events] == [
+        [1, 2, 3, 4],
+        [5, 6, 7, 8],
+        [1, 2, 3, 4],
+        [5, 6, 7, 8],
+    ]
+
+
 def test_take_events_opt_out_keeps_placeholders():
     tracker = _tracker(self_describing_kv_events=False)
     group_config = _group_config()

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

@@ -201,7 +201,7 @@ def test_swa_offload_window_covers_unaligned_hit(
     assert len(sched._jobs) == 1
     job = next(iter(sched._jobs.values()))
     if needs_left:
-        assert {left_key, swa[end - 1]} <= job.keys
+        assert {left_key, swa[end - 1]} <= set(job.keys)
     manager.complete_load(job.keys, state.req_context)
 
 
@@ -2615,7 +2615,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
     # First store: fires when block 0 is fully populated.
     runner.run(decoded_tokens=[0, 0], expected_stored=(0, 1, 2))
     assert runner.manager.complete_store.call_count == 1
-    first_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    first_call_keys = runner.manager.complete_store.call_args.args[0]
     assert len(first_call_keys) == 1
     runner.manager.complete_store.reset_mock()
 
@@ -2625,7 +2625,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
         expected_stored=(3, 4, 5),
     )
     assert runner.manager.complete_store.call_count == 1
-    second_call_keys = set(runner.manager.complete_store.call_args.args[0])
+    second_call_keys = runner.manager.complete_store.call_args.args[0]
     assert first_call_keys != second_call_keys
     runner.manager.complete_store.reset_mock()
 
@@ -2634,6 +2634,47 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
     assert runner.manager.complete_store.call_count == 0
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_store_job_preserves_order_and_deduplicates_keys(
+    request_runner, async_scheduling: bool
+):
+    tokens_per_block = 4
+    blocks_per_chunk = 3
+    tokens_per_chunk = tokens_per_block * blocks_per_chunk
+    runner = request_runner(
+        blocks_per_chunk=blocks_per_chunk,
+        block_size=tokens_per_block,
+        num_gpu_blocks=100,
+        async_scheduling=async_scheduling,
+    )
+    runner.new_request(token_ids=[0] * (tokens_per_chunk * 2))
+    returned_keys = []
+
+    def prepare_store(keys, req_context):
+        keys = list(keys)
+        returned_keys.extend(keys)
+        output = generate_store_output(keys)
+        output.keys_to_store.append(keys[0])
+        return output
+
+    runner.manager.prepare_store.side_effect = prepare_store
+    job_keys = {}
+
+    def capture_job_keys():
+        job_keys.update(runner.connector_scheduler._jobs.items())
+
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=tuple(range(6)),
+        post_step_fn=capture_job_keys,
+    )
+
+    assert len(job_keys) == 1
+    keys = next(iter(job_keys.values())).keys
+    assert keys == tuple(returned_keys)
+    assert runner.manager.complete_store.call_args.args[0] == keys
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_complete_store_waits_for_all_worker_acks(
     request_runner, async_scheduling: bool

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

@@ -95,6 +95,7 @@ class _OffloadEventMetadata:
     # Deferred: needs the same incremental curr_mm_idx handling as GPU events.
     extra_keys: tuple[tuple[Any, ...] | None, ...] | None
     group_idx: int
+    chunk_idx: int
     kv_cache_spec: OffloadingEventGroupSpec
     active_residencies: set[tuple[Medium, str | None]]
 
@@ -244,6 +245,7 @@ def _record_partial_tail(
             lora_name=lora_name,
             extra_keys=None,
             group_idx=group_config.group_idx,
+            chunk_idx=chunk_start // group_config.tokens_per_chunk,
             kv_cache_spec=group_config.kv_event_group_spec,
             active_residencies={_PRIMARY_RESIDENCY},
         )
@@ -345,6 +347,7 @@ def _build_event_metadata(
             lora_name=lora_name,
             extra_keys=None,
             group_idx=group_config.group_idx,
+            chunk_idx=chunk_idx,
             kv_cache_spec=group_config.kv_event_group_spec,
             active_residencies={_PRIMARY_RESIDENCY},
         )
@@ -384,9 +387,15 @@ def _placeholder_stored(
     def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]:
         # Metadata is read, NOT popped: the entry must survive until the
         # eviction event so BlockRemoved can fan out to the same hashes.
-        # Events are self-contained (own parent), so key order is free.
         locality = event.locality.value if event.locality is not None else None
-        for key in event.keys:
+
+        def key_order(key: OffloadKey) -> tuple[bool, int, int]:
+            meta = self._pending_event_metadata.get(key)
+            if meta is None:
+                return True, 0, 0
+            return False, meta.group_idx, meta.chunk_idx
+
+        for key in sorted(event.keys, key=key_order):
             meta = self._pending_event_metadata.get(key)
             if meta is None:
                 if self.self_describing_enabled:

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

@@ -77,8 +77,9 @@ class TransferJobStatus:
     # Number of workers still pending. Starts at num_workers,
     # decremented as each worker reports completion. Job is done at 0.
     pending_count: int
-    # Offload keys this job covers; passed to manager.complete_*().
-    keys: set[OffloadKey]
+    # Offload keys this job covers, in prefix order; passed to
+    # manager.complete_*().
+    keys: tuple[OffloadKey, ...]
     is_store: bool
     # Store source blocks fenced after the request finishes.
     deferred_fence_block_ids: list[int] | None = None
@@ -1356,7 +1357,7 @@ def update_state_after_alloc(
         self._jobs[load_job_id] = TransferJobStatus(
             req_id=request.request_id,
             pending_count=self.config.num_workers,
-            keys=set(keys_to_load),
+            keys=tuple(dict.fromkeys(keys_to_load)),
             is_store=False,
         )
 
@@ -1461,7 +1462,7 @@ def _build_aligned_boundary_store_jobs(
                 self._jobs[job_id] = TransferJobStatus(
                     req_id=req_id,
                     pending_count=self.config.num_workers,
-                    keys={key},
+                    keys=(key,),
                     is_store=True,
                     fenced_block_ids=[block_id],
                 )
@@ -1584,7 +1585,7 @@ def _build_partial_tail_store_jobs(
             self._jobs[job_id] = TransferJobStatus(
                 req_id=req_id,
                 pending_count=self.config.num_workers,
-                keys=set(store_output.keys_to_store),
+                keys=tuple(dict.fromkeys(store_output.keys_to_store)),
                 is_store=True,
                 fenced_block_ids=source_blocks,
             )
@@ -1829,7 +1830,8 @@ def _build_store_jobs(
                 req_status.advance_stored_idx(num_offloadable_tokens)
                 continue
 
-            keys_to_store = set(store_output.keys_to_store)
+            ordered_keys_to_store = tuple(dict.fromkeys(store_output.keys_to_store))
+            keys_to_store = set(ordered_keys_to_store)
 
             group_sizes: list[int] = []
             block_indices: list[int] = []
@@ -1902,7 +1904,7 @@ def _build_store_jobs(
             self._jobs[job_id] = TransferJobStatus(
                 req_id=req_id,
                 pending_count=self.config.num_workers,
-                keys=set(keys_to_store),
+                keys=ordered_keys_to_store,
                 is_store=True,
                 deferred_fence_block_ids=deferred_fence_block_ids,
                 fenced_block_ids=fenced_block_ids or None,