A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 00:00 UTC
Use CU_MEMCPY_SRC_ACCESS_ORDER_ANY for batch KV cache swaps (#39306)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Signed-off-by: Itay Etelis <etelis2019@gmail.com> Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Itay Etelis <etelis2019@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Itay Etelis committed
May 10, 2026 at 05:57 UTC
00b0618a03916f826e7da384a67ddf72f2302d97
5 files changed
+35
-6
csrc/cache.h
+2
-1
@@ -12,7 +12,8 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
12
13
void swap_blocks_batch(const torch::Tensor& src_ptrs,
14
const torch::Tensor& dst_ptrs,
15
- const torch::Tensor& sizes);
15
+ const torch::Tensor& sizes,
16
+ bool is_src_access_order_any);
17
18
void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
19
torch::Tensor& key_cache, torch::Tensor& value_cache,
csrc/cache_kernels.cu
+8
-2
@@ -77,7 +77,8 @@ void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
77
78
void swap_blocks_batch(const torch::Tensor& src_ptrs,
79
const torch::Tensor& dst_ptrs,
80
- const torch::Tensor& sizes) {
80
+ const torch::Tensor& sizes,
81
+ bool is_src_access_order_any) {
82
TORCH_CHECK(src_ptrs.device().is_cpu(), "src_ptrs must be on CPU");
83
TORCH_CHECK(dst_ptrs.device().is_cpu(), "dst_ptrs must be on CPU");
84
TORCH_CHECK(sizes.device().is_cpu(), "sizes must be on CPU");
@@ -124,7 +125,12 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
125
126
if (batch_fn != nullptr) {
127
CUmemcpyAttributes attr = {};
127
- attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
128
+ // ANY lets the DMA engine prefetch source bytes out of stream order,
129
+ // which is only safe when no GPU stream is concurrently writing the
130
+ // source.
131
+ attr.srcAccessOrder = is_src_access_order_any
132
+ ? CU_MEMCPY_SRC_ACCESS_ORDER_ANY
133
+ : CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
134
size_t attrs_idx = 0;
135
size_t fail_idx = 0;
136
CUresult result = batch_fn(reinterpret_cast<CUdeviceptr*>(dst_data),
csrc/torch_bindings.cpp
+2
-1
@@ -553,7 +553,8 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
553
// Batch swap: submit all block copies in a single driver call.
554
cache_ops.def(
555
"swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
556
- " Tensor sizes) -> ()");
556
+ " Tensor sizes,"
557
+ " bool is_src_access_order_any=False) -> ()");
558
cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
559
560
// Reshape the key and value tensors and cache them.
vllm/_custom_ops.py
+10
-1
@@ -2805,6 +2805,7 @@ def swap_blocks_batch(
2805
src_ptrs: torch.Tensor,
2806
dst_ptrs: torch.Tensor,
2807
sizes: torch.Tensor,
2808
+ is_src_access_order_any: bool = False,
2809
) -> None:
2810
"""
2811
Batch version of swap_blocks: submit all copies in a single driver call.
@@ -2813,8 +2814,16 @@ def swap_blocks_batch(
2814
of sizes[i] bytes. All three tensors must be int64 CPU tensors.
2815
On CUDA 12.8+ this uses cuMemcpyBatchAsync for minimal submission
2816
overhead; on older CUDA it falls back to a loop of cudaMemcpyAsync.
2817
+
2818
+ is_src_access_order_any: if True, pass CU_MEMCPY_SRC_ACCESS_ORDER_ANY to
2819
+ cuMemcpyBatchAsync, letting the DMA engine prefetch source bytes
2820
+ out of stream order. Only safe when no GPU stream is concurrently
2821
+ writing to the source. Defaults to False (STREAM ordering), which
2822
+ is always safe.
2823
"""
2817
- torch.ops._C_cache_ops.swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
2824
+ torch.ops._C_cache_ops.swap_blocks_batch(
2825
+ src_ptrs, dst_ptrs, sizes, is_src_access_order_any
2826
+ )
2827
2828
2829
def convert_fp8(
vllm/v1/kv_offload/cpu/gpu_worker.py
+13
-1
@@ -313,10 +313,22 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):
313
last_event = last_transfer.end_event
314
# assure job will start only after the previous one completes
315
stream.wait_event(last_event)
316
+ # CPU->GPU reads from host pinned memory, which is never written
317
+ # by a concurrent GPU stream, so CU_MEMCPY_SRC_ACCESS_ORDER_ANY is
318
+ # safe and lets the driver pipeline source reads. GPU->CPU reads
319
+ # from the live GPU KV cache, which the compute stream keeps
320
+ # writing; we must keep STREAM ordering so source reads are gated
321
+ # by the transfer stream's wait_stream(compute) barrier.
322
+ is_src_access_order_any = not self.gpu_to_cpu
323
with torch.cuda.stream(stream):
324
start_event.record(stream)
325
if num_copy_ops > 0:
319
- ops.swap_blocks_batch(batch_src, batch_dst, batch_sizes)
326
+ ops.swap_blocks_batch(
327
+ batch_src,
328
+ batch_dst,
329
+ batch_sizes,
330
+ is_src_access_order_any=is_src_access_order_any,
331
+ )
332
end_event.record(stream)
333
334
self._transfer_events[job_id] = end_event