[Bugfix][KV Offload] Fix swap_blocks_batch on the default stream (#46379)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <Itay.etelis@gmail.com>
Itay Etelis committed
Jun 23, 2026 at 15:45 UTC
7d47cff93380b7b20c903041b5204378f31ad758
2 files changed
+48
-1
csrc/libtorch_stable/cache_kernels.cu
+6
-1
@@ -127,7 +127,12 @@ void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
127
return reinterpret_cast<BatchFn>(fn_ptr);
128
}();
129
130
- if (batch_fn != nullptr) {
130
+ // cuMemcpyBatchAsync rejects the legacy default stream (handle 0 /
131
+ // cudaStreamLegacy) with CUDA_ERROR_INVALID_VALUE; route it to the per-copy
132
+ // fallback below, which is correct on any stream. Real and per-thread-default
133
+ // streams take the batch fast path.
134
+ const bool usable_stream = stream != nullptr && stream != cudaStreamLegacy;
135
+ if (batch_fn != nullptr && usable_stream) {
136
CUmemcpyAttributes attr = {};
137
// ANY lets the DMA engine prefetch source bytes out of stream order,
138
// which is only safe when no GPU stream is concurrently writing the
tests/v1/kv_offload/cpu/test_swap_blocks_batch.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the ``ops.swap_blocks_batch`` C++ (cuMemcpyBatchAsync) path."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.platforms import current_platform
10
+
11
+
12
+def _addrs(buffers: list[torch.Tensor]) -> torch.Tensor:
13
+ return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64)
14
+
15
+
16
+def _run_batch(sizes: list[int]) -> None:
17
+ src = [torch.randint(256, (s,), dtype=torch.uint8, device="cuda") for s in sizes]
18
+ dst = [torch.zeros_like(s) for s in src]
19
+ ops.swap_blocks_batch(
20
+ _addrs(src), _addrs(dst), torch.tensor(sizes, dtype=torch.int64)
21
+ )
22
+ torch.accelerator.synchronize()
23
+ for s, d in zip(src, dst):
24
+ assert torch.equal(d, s)
25
+
26
+
27
+@pytest.mark.skipif(
28
+ not current_platform.is_cuda(), reason="swap_blocks_batch requires CUDA"
29
+)
30
+def test_swap_blocks_batch_default_stream():
31
+ # cuMemcpyBatchAsync rejects the legacy default stream; the op must fall
32
+ # back to per-copy transfers instead of raising.
33
+ _run_batch([8, 4096, 8192])
34
+
35
+
36
+@pytest.mark.skipif(
37
+ not current_platform.is_cuda(), reason="swap_blocks_batch requires CUDA"
38
+)
39
+def test_swap_blocks_batch_dedicated_stream():
40
+ # A dedicated non-default stream exercises the cuMemcpyBatchAsync fast path.
41
+ with torch.cuda.stream(torch.cuda.Stream()):
42
+ _run_batch([8, 4096, 8192])