[KV offload] Order CPU->GPU loads against the compute stream (#50696)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Itay Etelis committed
Aug 31, 2026 at 14:31 UTC
bed3280f50dce36ff9472eee937d1aada7e31d75
1 file changed
+4
-2
vllm/v1/kv_offload/cpu/gpu_worker.py
+4
-2
@@ -636,8 +636,10 @@ class SingleDirectionOffloadingHandler:
636
637
# Stores must wait for the model to finish writing the KV they read.
638
# Loads must wait for pending writes (including zeroing) to their
639
- # destination blocks; otherwise an earlier transfer can be overwritten
640
- # by compute-stream work that was already queued when the load began.
639
+ # destination blocks: the scheduler's _skip_zero_block_ids only edits
640
+ # the step being scheduled and cannot retract zeroing shipped in an
641
+ # earlier step for a since-reallocated block; with async scheduling
642
+ # nothing else orders that zeroing against this copy.
643
stream.wait_stream(current_platform.current_stream())
644
if self._transfers:
645
last_transfer: Transfer = self._transfers[-1]