[Bugfix][KV Offload] Order SimpleCPUOffloadConnector stores behind the compute stream

[Peuqui]

Purpose

SimpleCPUOffloadWorker.get_finished() issues the GPU→CPU copy on its own transfer stream, from a background thread, without ever ordering that stream behind the compute stream. On top of that, build_params() requests CU_MEMCPY_SRC_ACCESS_ORDER_ANY for both directions.

The store therefore reads the live KV cache while the compute stream is still writing into it, and the driver is explicitly permitted to pull those source reads ahead of anything else on the stream. A block can be captured half-written. Its hash is stamped regardless, so the damaged copy enters the CPU prefix cache as valid data. Nothing fails and nothing is logged; a later request whose prefix matches restores the block and answers from a poisoned context.

The sibling path in this repository already handles this correctly. vllm/v1/kv_offload/cpu/gpu_worker.py does stream.wait_stream(torch.cuda.current_stream()) before offloading and deliberately keeps STREAM source ordering for GPU→CPU, with a comment stating why: the compute stream keeps writing the source. SimpleCPUOffloadConnector does neither.

Both halves are needed. The barrier alone is not sufficient, because ANY lets the driver read past it. ANY stays for CPU→GPU, whose source is host pinned memory that no GPU stream writes, so it can keep pipelining those reads.

How it was found

Measured on DeepSeek-V4-Flash, 5-stage pipeline parallel (2× RTX 8000 + 3× V100), FP8 KV cache, cpu_bytes_to_use 5 GiB. Request sequence: two 18k requests, a 30k needle-in-a-haystack, then a 62k needle that shares a ~2k-token prefix with the 30k one, so its opening blocks are restored from the CPU pool.

Each completed store was re-read and compared against its GPU source, right after the transfer event completed. The scheduler still holds a reference on those GPU blocks until every worker reports the event, so the source cannot have been reused in the meantime — any difference means the copy captured data the compute stream had not finished writing.

Without this change:

PP1: 21504 block/tensor pairs compared, 12 differ
PP2: 21504 compared,  9 differ
PP3: 21504 compared,  8 differ
PP0, PP4: 0 differ

event=0  model.layers.17.attn                 gpu_block=28  cpu_block=1     583 of  1728 bytes differ
event=0  model.layers.18.attn.indexer.k_cache gpu_block=28  cpu_block=1     129 of  8640 bytes differ
event=89 model.layers.16.attn                 gpu_block=598 cpu_block=2489 2255 of 37440 bytes differ

Partial differences in the middle of a block — a mix of old and new content, which is what a copy looks like when it races the writer. End to end, the 62k needle returned 0 of 4 facts and degenerated into token salad.

With this change, same sequence, same comparison: zero differing pairs on all five ranks, needle 4 of 4, repeated twice. Step time is unchanged, 82/84 ms against 81/83 ms before, so the barrier costs nothing measurable on this setup.

Please note the fault is timing-dependent: the identical sequence produced both a clean 4-of-4 run and a corrupted one under otherwise identical conditions. Judging this by output quality alone is unreliable, which is why the block comparison was used. A regression test that reliably fails without the fix therefore cannot exist; the added tests pin the structure (barrier present, source ordering correct per direction), not the timing.

Test Plan

pytest tests/v1/simple_kv_offload/
pre-commit run --files vllm/v1/simple_kv_offload/worker.py \
    vllm/v1/simple_kv_offload/copy_backend.py \
    vllm/v1/simple_kv_offload/cuda_mem_ops.py \
    tests/v1/simple_kv_offload/test_store_stream_ordering.py

New tests/v1/simple_kv_offload/test_store_stream_ordering.py covers: the store queues the barrier before the copy is handed to the backend, a load does not take one, and build_params selects STREAM ordering for stores and ANY for loads.

End-to-end on the hardware above: the sequence described in Purpose, run twice, plus a restore probe (30k needle, 62k evictor with a different seed, identical 30k needle again) verifying the answer is unchanged.

Test Result

$ pytest tests/v1/simple_kv_offload/
19 passed, 8 skipped, 14 warnings in 14.07s

$ pre-commit run --files <the four files above>
ruff check ... Passed
ruff format ... Passed
typos ... Passed
Run mypy locally for lowest supported Python version ... Passed
Check SPDX headers ... Passed
Prevent new 'torch.cuda' APIs call ... Passed
(all other hooks Passed or Skipped)

Against the unpatched sources, test_store_is_ordered_behind_the_compute_stream fails with assert 'store_stream.wait_stream' in ['backend.launch_copy'], i.e. the copy is queued with no ordering at all.

End-to-end: zero differing block/tensor pairs, needle 4 of 4 in both runs, restore probe returns a character-identical answer, step time unchanged.

Duplicate check

• gh pr list --repo 1CatAI/1Cat-vLLM --state open --search "offload" — [KV Offload] Grouped hybrid RAM caching, filesystem restoration and retention-sized Mamba pools (stacked on #617) #598 and [Bugfix][Core] Backport upstream Mamba prefix-cache retention #617 touch the native/QSA offload path, [Bugfix] PLE offload: send registration inputs by file descriptor #654 and [Feature][Qwen4Exp] PLE overflow cascade: device, pinned host, spare GPU, disk #646 the PLE offload, none of them SimpleCPUOffloadConnector.
• gh pr list --state all --search "SimpleCPUOffloadConnector in:body" — only [Bugfix][KV Offloading] Allow native CPU offload with expandable_segments #248 (merged, expandable_segments), a different fix.
• gh pr list --state all --search "wait_stream srcAccessOrder" — no results.

AI assistance

This change was developed with AI assistance (Claude). The diagnosis, the instrumentation used to produce the numbers above, and the measurements were carried out on the hardware described. Every changed line has been reviewed by me and I can defend the change end to end.