[Perf] Batch KV cache swap copies via cuMemcpyBatchAsync

[Etelis]

Replace per-layer per-block swap_blocks calls in the KV cache offloading
handler with a single swap_blocks_batch call that submits all copies in
one driver invocation.

On CUDA 12.8+ this uses cuMemcpyBatchAsync; on older CUDA/ROCm it falls
back to a flat cudaMemcpyAsync loop with cudaMemcpyDefault. Zero extra
GPU memory. No behavior change.

Supersedes #38216 (rebased clean).

Benchmark Results

Hardware: 8xH100 80GB HBM3, CUDA 12.8/12.9.

Baseline runs the original files via
pip install vllm==0.18.0 --force-reinstall --no-deps. Each mode ran in a
fresh Python process to avoid stale module state.

Handler-level benchmark

Directly measures the transfer path in isolation — no model inference.

Setup: Instantiate CpuGpuOffloadingHandlers with BF16 GPU tensors shaped
by FlashAttentionBackend.get_kv_cache_shape() using each model's real
architecture.
Call handler.transfer_async(), poll get_finished() until complete. Measure
with time.perf_counter() (includes CUDA sync). 5 warmup + 100 measured
iterations. Baseline and batched ran as separate processes.

Config	Layers	KV Heads	Head Dim	Tensors	Per-tensor block size
LLaMA-8B	32	8	128	64	32 KB
LLaMA-70B	80	8	128	160	32 KB
LLaMA-70B TP=4	80	2	128	160	8 KB
Qwen2.5-0.5B	24	2	64	48	4 KB
Qwen2.5-1.5B	28	2	128	56	8 KB
Qwen2.5-3B	36	2	128	72	8 KB
Phi-3-mini	32	8	96	64	24 KB

Per-layer tensors (FlashAttn K/V split, no cross-layer allocation):

Config	Tensors	Baseline	Batched	Speedup
LLaMA-8B, 64 blocks	64	16,358 us	4,519 us	3.6x
LLaMA-70B, 64 blocks	160	41,124 us	10,239 us	4.0x
LLaMA-70B TP=4, 32 blocks	160	20,427 us	3,299 us	6.2x
Qwen2.5-0.5B, 128 blocks	48	23,484 us	3,184 us	7.4x
Qwen2.5-1.5B, 64 blocks	56	14,042 us	2,272 us	6.2x
Qwen2.5-3B, 64 blocks	72	18,045 us	2,935 us	6.2x
Phi-3-mini, 64 blocks	64	15,963 us	3,888 us	4.1x

E2E vLLM serve — KV transfer bandwidth

Setup:

1. Start vllm serve with each real model, offloading enabled:

   python -m vllm.entrypoints.openai.api_server \
       --model <model> --gpu-memory-utilization <0.4-0.5> \
       --kv-transfer-config '{\"kv_connector\":\"OffloadingConnector\",
           \"kv_role\":\"kv_both\",
           \"kv_connector_extra_config\":{\"cpu_bytes_to_use\":2147483648,
           \"block_size\":48}}' \
       --max-model-len 4096 --enforce-eager

2. Send sustained load: 8 concurrent Python threads, each sending prompts in a loop for 45 seconds.
3. Read the KV transfer bandwidth from vLLM's server log
   total_time is measured by CUDA events (start_event.elapsed_time(end_event))
   inside SingleDirectionOffloadingHandler.get_finished() Bandwidth = total_bytes / total_time.

Model	KV Heads	Baseline BW	Batched BW	Improvement
LLaMA-3.2-1B	8	27.9 GB/s	32.4 GB/s	+16%
Qwen2.5-1.5B	2	25.8 GB/s	31.8 GB/s	+23%
Qwen2.5-3B	2	29.5 GB/s	34.3 GB/s	+16%
Gemma-2-2B	4	40.2 GB/s	43.3 GB/s	+8%
Phi-3.5-mini	8	51.6 GB/s	53.3 GB/s	+3%
Qwen2.5-7B	4	34.2 GB/s	40.7 GB/s	+19%
LLaMA-3.1-8B	8	44.0 GB/s	47.0 GB/s	+7%
Mistral-7B	8	45.1 GB/s	48.6 GB/s	+8%

3. E2E serving throughput

Model	Baseline req/s	Batched req/s	Baseline p99 TTFT	Batched p99 TTFT
Qwen2.5-7B	6.0	7.1 (+18%)	6,077 ms	117 ms
LLaMA-3.1-8B	6.7	6.7	100.7 ms	105.7 ms
Mistral-7B	6.7	6.6	91.4 ms	82.5 ms

Qwen2.5-7B shows 18% throughput improvement and p99 TTFT reduction from
6 seconds to 117ms under concurrent load. This I guess it because the model has 4 KV heads
(vs 8 for LLaMA/Mistral), producing smaller per-block copies where
submission overhead dominates.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request introduces swap_blocks_batch, a batched memory copy operation designed to reduce overhead during KV cache offloading. The implementation leverages cuMemcpyBatchAsync on supported CUDA versions (12.8+) and provides a fallback for older environments. The CPUGPUWorker is updated to aggregate transfer requests into a single batch call. A review comment suggests that the operation should be registered for the CUDA device in the Torch bindings instead of CPU.

[gemini-code-assist]

The swap_blocks_batch operation is implemented in csrc/cache_kernels.cu and involves CUDA operations. However, it's being registered here for the torch::kCPU device. This is likely incorrect as the implementation relies on CUDA streams and memory copies. It should be registered for torch::kCUDA to ensure it's dispatched correctly when called on CUDA tensors.

Suggested change

	cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
	cache_ops.impl("swap_blocks_batch", torch::kCUDA, &swap_blocks_batch);

[Etelis]

The input tensors (src_ptrs, dst_ptrs, sizes) are CPU tensors — they're numpy arrays of raw pointers/sizes converted via torch.from_numpy(). PyTorch dispatches based on the input tensor device, so kCPU is correct here. The existing swap_blocks uses kCUDA because its inputs are the actual GPU KV cache tensors. Registering with kCUDA would actually break dispatch since no input tensor lives on GPU.

[ivanium]

Minor comment: Curious have you tried CU_MEMCPY_SRC_ACCESS_ORDER_ANY (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f)? I found it gives me better CPU->GPU bandwidth on Grace Blackwell nodes.

[orozery]

I see you also applied this parameter to GPU srcs.
According to the documentation this means access to srcs can be out of stream, so potentially not waiting for the compute (default) stream to complete?

@Etelis Anyhow for CPU->GPU this seems safe. Let's test it towards a follow up.

Thanks @ivanium for this suggestion!

[ivanium]

Right, it won't wait for previous ops in the stream. Since we typically call this API in a separate copy stream, I guess we cannot rely on this AccessOrder param anyway. If we want to stay safe as a general purpose API, maybe we can expose a configurable param to users.

[Etelis]

I'll test it as a followup