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 14, 05:30 UTC
[Nixl][P/D] Add cuda2cpu support (HD->DH transfer) (#24690)
Signed-off-by: Chenxi Yang <cxyang@fb.com> Co-authored-by: Chenxi Yang <cxyang@fb.com>
Chenxi Yang committed
Sep 29, 2025 at 07:31 UTC
d0d138bc55f707e86c66b1d4edfb3f4046e9aa5f
6 files changed
+96
-15
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+28
-3
@@ -1,6 +1,31 @@
1
#!/bin/bash
2
set -xe
3
4
+# Parse command line arguments
5
+KV_BUFFER_DEVICE="cuda" # Default to cuda
6
+while [[ $# -gt 0 ]]; do
7
+ case $1 in
8
+ --kv_buffer_device)
9
+ KV_BUFFER_DEVICE="$2"
10
+ shift 2
11
+ ;;
12
+ *)
13
+ echo "Unknown option $1"
14
+ echo "Usage: $0 [--kv_buffer_device <cuda|cpu>]"
15
+ exit 1
16
+ ;;
17
+ esac
18
+done
19
+
20
+echo "Running accuracy tests with kv_buffer_device=$KV_BUFFER_DEVICE"
21
+
22
+# Build the kv-transfer-config once
23
+if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
24
+ KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
25
+else
26
+ KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}"
27
+fi
28
+
29
# Models to run
30
MODELS=(
31
"Qwen/Qwen3-0.6B"
@@ -79,7 +104,7 @@ run_tests_for_model() {
104
105
# Calculate port number (base port + instance number)
106
PORT=$((8100 + i))
82
- # Calculate side channel port. Avoid clash with with TP workers.
107
+ # Calculate side channel port. Avoid clash with with TP workers.
108
SIDE_CHANNEL_PORT=$((5559 + i))
109
110
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
@@ -93,7 +118,7 @@ run_tests_for_model() {
118
--enforce-eager \
119
--gpu-memory-utilization 0.2 \
120
--tensor-parallel-size $PREFILLER_TP_SIZE \
96
- --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
121
+ --kv-transfer-config '$KV_CONFIG'"
122
123
if [ -n "$model_args" ]; then
124
FULL_CMD="$BASE_CMD $model_args"
@@ -128,7 +153,7 @@ run_tests_for_model() {
153
--enforce-eager \
154
--gpu-memory-utilization 0.2 \
155
--tensor-parallel-size $DECODER_TP_SIZE \
131
- --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
156
+ --kv-transfer-config '$KV_CONFIG'"
157
158
if [ -n "$model_args" ]; then
159
FULL_CMD="$BASE_CMD $model_args"
tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh
+32
-5
@@ -1,6 +1,33 @@
1
#!/bin/bash
2
set -xe
3
4
+# Parse command line arguments
5
+KV_BUFFER_DEVICE="cuda" # Default to cuda
6
+PREFILL_GPU_ID=4 # Default GPU IDs
7
+DECODE_GPU_ID=5
8
+while [[ $# -gt 0 ]]; do
9
+ case $1 in
10
+ --kv_buffer_device)
11
+ KV_BUFFER_DEVICE="$2"
12
+ shift 2
13
+ ;;
14
+ *)
15
+ echo "Unknown option $1"
16
+ echo "Usage: $0 [--kv_buffer_device <cuda|cpu>]"
17
+ exit 1
18
+ ;;
19
+ esac
20
+done
21
+
22
+echo "Running edge case tests with kv_buffer_device=$KV_BUFFER_DEVICE (GPUs: $PREFILL_GPU_ID, $DECODE_GPU_ID)"
23
+
24
+# Build the kv-transfer-config once
25
+if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
26
+ KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
27
+else
28
+ KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\"}"
29
+fi
30
+
31
# Models to run
32
MODELS=(
33
"Qwen/Qwen3-0.6B"
@@ -50,15 +77,15 @@ run_tests_for_model() {
77
78
# Get model-specific arguments
79
local model_args=$(get_model_args "$model_name")
53
-
80
+
81
# Start prefill instance
82
PREFILL_PORT=8001
83
57
- BASE_CMD="CUDA_VISIBLE_DEVICES=0 VLLM_NIXL_SIDE_CHANNEL_PORT=5559 vllm serve $model_name \
84
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID VLLM_NIXL_SIDE_CHANNEL_PORT=5559 vllm serve $model_name \
85
--port $PREFILL_PORT \
86
--enforce-eager \
87
--gpu-memory-utilization 0.2 \
61
- --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
88
+ --kv-transfer-config '$KV_CONFIG'"
89
90
if [ -n "$model_args" ]; then
91
FULL_CMD="$BASE_CMD $model_args"
@@ -72,11 +99,11 @@ run_tests_for_model() {
99
DECODE_PORT=8002
100
101
# Build the command with or without model-specific args
75
- BASE_CMD="CUDA_VISIBLE_DEVICES=1 VLLM_NIXL_SIDE_CHANNEL_PORT=6000 vllm serve $model_name \
102
+ BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU_ID VLLM_NIXL_SIDE_CHANNEL_PORT=6000 vllm serve $model_name \
103
--port $DECODE_PORT \
104
--enforce-eager \
105
--gpu-memory-utilization 0.2 \
79
- --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"
106
+ --kv-transfer-config '$KV_CONFIG'"
107
108
if [ -n "$model_args" ]; then
109
FULL_CMD="$BASE_CMD $model_args"
vllm/config/kv_transfer.py
+2
-2
@@ -28,8 +28,8 @@ class KVTransferConfig:
28
"""The engine id for KV transfers."""
29
30
kv_buffer_device: Optional[str] = "cuda"
31
- """The device used by kv connector to buffer the KV cache.
32
- Currently only support 'cuda'."""
31
+ """The device used by kv connector to buffer the KV cache. Choices are
32
+ 'cuda' and 'cpu'."""
33
34
kv_buffer_size: float = 1e9
35
"""The buffer size for TorchDistributedConnector. Measured in number of
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+7
-1
@@ -67,7 +67,10 @@ except ImportError:
67
# Supported platforms and types of kv transfer buffer.
68
# {device: tuple of supported kv buffer types}
69
_NIXL_SUPPORTED_DEVICE = {
70
- "cuda": ("cuda", ),
70
+ "cuda": (
71
+ "cuda",
72
+ "cpu",
73
+ ),
74
"tpu": ("cpu", ),
75
"xpu": ("cpu", ),
76
}
@@ -701,6 +704,9 @@ class NixlConnectorWorker:
704
705
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
706
"""Assign copy (d2h, h2d) operations when host buffer is used."""
707
+ # Set a no-op if the host buffer is not cpu.
708
+ if self.kv_buffer_device != "cpu":
709
+ return
710
assert self.use_host_buffer
711
self.copy_blocks = copy_operation
712
vllm/platforms/cuda.py
+24
@@ -500,6 +500,30 @@ class CudaPlatformBase(Platform):
500
"You can use float16 instead by explicitly setting the "
501
"`dtype` flag in CLI, for example: --dtype=half.")
502
503
+ @classmethod
504
+ def insert_blocks_to_device(
505
+ cls,
506
+ src_cache: torch.Tensor,
507
+ dst_cache: torch.Tensor,
508
+ src_block_indices: torch.Tensor,
509
+ dst_block_indices: torch.Tensor,
510
+ ) -> None:
511
+ """Copy blocks from src_cache to dst_cache on GPU."""
512
+ _src_cache = src_cache[:, src_block_indices]
513
+ dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
514
+
515
+ @classmethod
516
+ def swap_out_blocks_to_host(
517
+ cls,
518
+ src_cache: torch.Tensor,
519
+ dst_cache: torch.Tensor,
520
+ src_block_indices: torch.Tensor,
521
+ dst_block_indices: torch.Tensor,
522
+ ) -> None:
523
+ """Copy blocks from GPU to host (CPU)."""
524
+ _src_cache = src_cache[:, src_block_indices]
525
+ dst_cache[:, dst_block_indices] = _src_cache.cpu()
526
+
527
@classmethod
528
def support_hybrid_kv_cache(cls) -> bool:
529
return True
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -4059,10 +4059,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4059
self.drafter.validate_same_kv_cache_group(kv_cache_config)
4060
4061
if has_kv_transfer_group():
4062
- get_kv_transfer_group().register_kv_caches(kv_caches)
4063
- if self.device.type == 'xpu':
4064
- get_kv_transfer_group().set_host_xfer_buffer_ops(
4065
- copy_kv_blocks)
4062
+ kv_transfer_group = get_kv_transfer_group()
4063
+ kv_transfer_group.register_kv_caches(kv_caches)
4064
+ kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
4065
4066
if self.dcp_world_size > 1:
4067
layer_names = self.attn_groups[0][0].layer_names