[Bugfix][KVConnector] Fix SimpleCPUOffloadConnector GPU->CPU store race

[Saddss]

Purpose

Fixes #45704.

SimpleCPUOffloadConnector defers both loads and stores to get_finished(), which runs after model execution. It assumes that blocks whose tokens are host-side "confirmed" by the scheduler are also safe to read on the device, and therefore launches the GPU->CPU store with no cross-stream synchronization.

That assumption does not hold under v1 overlapped execution. The host runs ahead of the GPU, so the KV-write kernels for the blocks being stored may still be in flight on the compute stream when the store is launched on a separate copy stream. The copy backend additionally uses srcAccessOrder=ANY for every transfer, which allows the DMA engine to read the source before prior writes to it have completed.

The combination lets the store read partially written or stale KV blocks and silently corrupt the CPU cache. There is no error or crash; the only symptom is garbled output on later requests that hit the corrupted CPU blocks.

This is the same class of hazard that #31341 fixed for OffloadingConnector; SimpleCPUOffloadConnector reintroduced it because it dropped the cross-stream barrier.

Fix

• Record a compute-done event on the current (compute) stream in get_finished() and have the copy backend wait on it before issuing the GPU->CPU store, so the store is ordered after the kernels that wrote those KV blocks.
• Use srcAccessOrder=STREAM for stores (the source is the live KV cache) and keep srcAccessOrder=ANY for loads (the source is stable pinned host memory).
• Loads are unchanged: they still launch immediately with no barrier, since they read from stable pinned host memory.

The fix preserves the connector's existing design (precise cudaHostRegister pinning, batched DMA, background copy thread) and only adds the missing happens-before edge for stores.

Test Plan

• New unit test tests/v1/simple_kv_offload/test_worker.py:
  • test_store_orders_after_compute_write: drives DmaCopyBackend directly and uses torch.cuda._sleep to deterministically delay the KV write. Asserts both directions so it is self-validating: the no-barrier control must corrupt (proving the race window is exercised) and the fixed path must be clean.
  • test_get_finished_passes_wait_event_for_store_only: asserts get_finished() passes a compute-done event for stores and None for loads.
  • test_build_params_src_access_order: asserts build_params defaults to ANY and honors an explicit STREAM override.
• Command: pytest tests/v1/simple_kv_offload/test_worker.py -v

Test Result

• Unit tests: 3 passed on the patched build. The same test file fails to even import against an unpatched build (the new symbols do not exist), confirming it covers the change.
• Deterministic race microbenchmark over 150 iterations: without the barrier the store corrupts 150/150 iterations; with the barrier 0/150.
• End-to-end with the production workload that originally surfaced [Bug]: SimpleCPUOffloadConnector corrupts restored KV under load → intermittent garbled output (missing GPU→CPU store cross-stream sync) #45704: garbled output drops from a steady non-zero rate to 0 with the fix applied.
• Performance A/B with CPU offload enabled (same workload): output throughput and mean TPOT are within run-to-run noise (output throughput -0.47%, mean TPOT -0.1%), i.e. no measurable regression.

[aoshen02]

Hi, thanks for the contribution, I wonder do you guys ever try using mooncake store connector?

[Saddss]

Hi, thanks for the contribution, I wonder do you guys ever try using mooncake store connector?嗨，感谢您的贡献，我想知道你们有没有尝试过使用 mooncake store connector？

Thanks for the pointer! Yes, we're aware of MooncakeStoreConnector and it's a great piece of work — for setups with RDMA it clearly offers more (cross-instance prefix sharing, SSD tiering, GPUDirect transfers).

Our deployment is a bit of an awkward fit for it, though: the fleet is all RTX 5090 with no RDMA NICs. Mooncake's main advantage is the RDMA / GPUDirect path, and without an RDMA HCA the transfer engine falls back to TCP. We did test Mooncake as a shared DRAM pool, but over TCP the transfer bandwidth couldn't keep up with recompute — a KV cache hit ended up slower than just recomputing the prefix, so it was a net loss for us.

For a single-node, no-RDMA setup, SimpleCPUOffloadConnector fits well: it does a direct cuMemcpyBatchAsync between GPU HBM and local pinned host memory over PCIe, with no network stack in the path, so it saturates PCIe bandwidth.

[Saddss]

Contributor

Another reason we depend on SimpleCPUOffloadConnector specifically is precise host memory allocation. Our nodes have a hard whole-machine memory budget, so the CPU KV pool has to use exactly what we configure in order to maximize usable cache.

The legacy CPU offload path can't give us that, because of PyTorch's pinned caching allocator. It allocates the host pool as one torch.zeros(..., pin_memory=True) per KV cache tensor (vllm/v1/kv_offload/cpu/gpu_worker.py), and that allocator rounds every allocation up to the next power of two (roundSize = c10::llvm::PowerOf2Ceil(size) in aten/src/ATen/core/CachingHostAllocator.h; see also pytorch/pytorch#150517). Since the rounding is per tensor, the overshoot compounds. For our gemma-4 setup the pool is split into 5 CPU tensors: budgeting 60 GB gives ~12 GB per tensor, and PowerOf2Ceil(12 GB) = 16 GB, so the host actually allocates 5 × 16 = 80 GB. In the worst case (each tensor just above a power-of-two boundary) this approaches a 2x blow-up — under a fixed memory ceiling that either wastes a large fraction of the cache budget or pushes us into OOM.

SimpleCPUOffloadConnector avoids this by construction: it allocates plain (unpinned) memory and then pins it with cudaHostRegister, bypassing the caching host allocator entirely, so the pool is exactly the size we configure (55 GB stays 55 GB, not 64 GB). For a tight memory budget that exact sizing is a concrete operational reason we rely on this connector — and why the correctness fix in this PR matters for us.

[ivanium]

Thanks for the fix! Overall good; left a few comments

[ivanium]

Thanks for the effort! LGTM in general and I'd approve. I do suggest making another pass of comments to be more concise because agents tend to comment too much. I commented on two major ones but some other comment blocks worth double check too. Otherwise all good.

[Saddss]

Thanks for the effort! LGTM in general and I'd approve. I do suggest making another pass of comments to be more concise because agents tend to comment too much. I commented on two major ones but some other comment blocks worth double check too. Otherwise all good.

You're right, and thanks for calling it out. To be transparent about our workflow: we scope the problem and decide the fix direction, then have an agent drive the actual implementation and testing. It moves fast, but over-commenting is one of the tells, as you spotted. Did another pass: removed the gate comment in the copy loop, de-duplicated the STREAM/ANY rationale so it lives in a single place (semantics on the enum, the per-direction choice at the call site), and tightened the rest. Going forward I'll keep comments lean from the start instead of leaning on a review round to trim them. Appreciate the patient review.

[ivanium]

Thanks for the effort! LGTM in general and I'd approve. I do suggest making another pass of comments to be more concise because agents tend to comment too much. I commented on two major ones but some other comment blocks worth double check too. Otherwise all good.

You're right, and thanks for calling it out. To be transparent about our workflow: we scope the problem and decide the fix direction, then have an agent drive the actual implementation and testing. It moves fast, but over-commenting is one of the tells, as you spotted. Did another pass: removed the gate comment in the copy loop, de-duplicated the STREAM/ANY rationale so it lives in a single place (semantics on the enum, the per-direction choice at the call site), and tightened the rest. Going forward I'll keep comments lean from the start instead of leaning on a review round to trim them. Appreciate the patient review.

No worries and thanks for the nice fix! I have enabled CI and it should be merged once CI passes which I don't anticipate any issues

[Saddss]

Thanks for the effort! LGTM in general and I'd approve. I do suggest making another pass of comments to be more concise because agents tend to comment too much. I commented on two major ones but some other comment blocks worth double check too. Otherwise all good.

You're right, and thanks for calling it out. To be transparent about our workflow: we scope the problem and decide the fix direction, then have an agent drive the actual implementation and testing. It moves fast, but over-commenting is one of the tells, as you spotted. Did another pass: removed the gate comment in the copy loop, de-duplicated the STREAM/ANY rationale so it lives in a single place (semantics on the enum, the per-direction choice at the call site), and tightened the rest. Going forward I'll keep comments lean from the start instead of leaning on a review round to trim them. Appreciate the patient review.

No worries and thanks for the nice fix! I have enabled CI and it should be merged once CI passes which I don't anticipate any issues

There has only two failures, and both look unrelated to this PR:
e2e-core-1-gpu: test_cascade_attention[FLASH_ATTN] — minor output string mismatch
metrics-tracing-2-gpus: test_traces — SIGSEGV after generation finished
Could you retry those two failed jobs? @ivanium

[Saddss]

Thanks for the effort! LGTM in general and I'd approve. I do suggest making another pass of comments to be more concise because agents tend to comment too much. I commented on two major ones but some other comment blocks worth double check too. Otherwise all good.

You're right, and thanks for calling it out. To be transparent about our workflow: we scope the problem and decide the fix direction, then have an agent drive the actual implementation and testing. It moves fast, but over-commenting is one of the tells, as you spotted. Did another pass: removed the gate comment in the copy loop, de-duplicated the STREAM/ANY rationale so it lives in a single place (semantics on the enum, the per-direction choice at the call site), and tightened the rest. Going forward I'll keep comments lean from the start instead of leaning on a review round to trim them. Appreciate the patient review.

No worries and thanks for the nice fix! I have enabled CI and it should be merged once CI passes which I don't anticipate any issues

There has only two failures, and both look unrelated to this PR: e2e-core-1-gpu: test_cascade_attention[FLASH_ATTN] — minor output string mismatch metrics-tracing-2-gpus: test_traces — SIGSEGV after generation finished Could you retry those two failed jobs? @ivanium

CI is all green on. Could you help get it merged？Thanks😊