Skip to content

Commit ba930de

Browse files
MC952-archleoda1
MC952-arch
and
leoda1
authored
[Others] KV transfer benchmark (#474)
--------- Co-authored-by: Da Liu <liud33787@gmail.com>
1 parent e774d8a commit ba930de

3 files changed

Lines changed: 1273 additions & 0 deletions

File tree

‎test/perf/kv_transfer/README.md‎

Lines changed: 295 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,295 @@
1+
# KV Transfer Connector Benchmark
2+
3+
Unified benchmark for measuring raw transport-level performance of vLLM KV transfer connectors (NIXL, Mooncake, and FlagCX).
4+
5+
## Overview
6+
7+
This benchmark measures bandwidth and latency for KV cache transfers using the same underlying transport libraries as vLLM's KV connectors, without the full vLLM scheduler/worker overhead.
8+
9+
**Supported Connectors:**
10+
- **NIXL** with UCX or FLAGCX backends
11+
- **Mooncake** TransferEngine (RDMA)
12+
- **FlagCX** direct one-sided RDMA (no NIXL abstraction layer)
13+
14+
## Requirements
15+
16+
### NIXL
17+
```bash
18+
pip install nixl-cu12 # or nixl-cu11 depending on CUDA version
19+
pip install pyzmq torch
20+
```
21+
22+
For FLAGCX backend:
23+
```bash
24+
# Ensure libflagcx.so is in LD_LIBRARY_PATH
25+
# Ensure libplugin_FLAGCX.so is on the NIXL plugin path
26+
```
27+
28+
### Mooncake
29+
```bash
30+
# Install Mooncake following:
31+
# https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md
32+
pip install pyzmq torch
33+
```
34+
35+
### FlagCX (Direct)
36+
```bash
37+
# Build FlagCX from source:
38+
# https://github.com/FlagOpen/FlagCX
39+
# Set FLAGCX_PATH to the FlagCX root directory
40+
export FLAGCX_PATH=/path/to/FlagCX
41+
pip install pyzmq torch
42+
```
43+
44+
## Usage
45+
46+
### NIXL with UCX Backend
47+
48+
**Server (Node A, IP: 10.8.2.169):**
49+
```bash
50+
python kv_transfer_benchmark.py --connector=nixl --role=server \
51+
--remote-ip=10.8.2.169 --device=gpu --nixl-backend=UCX \
52+
--sizes=1048576,16777216,104857600 --iters=20
53+
```
54+
55+
**Client (Node B):**
56+
```bash
57+
python kv_transfer_benchmark.py --connector=nixl --role=client \
58+
--remote-ip=10.8.2.169 --device=gpu --nixl-backend=UCX \
59+
--sizes=1048576,16777216,104857600 --iters=20
60+
```
61+
62+
### NIXL with FLAGCX Backend
63+
64+
Same as above, but add `--nixl-backend=FLAGCX`:
65+
```bash
66+
python kv_transfer_benchmark.py --connector=nixl --role=server \
67+
--remote-ip=10.8.2.169 --device=gpu --nixl-backend=FLAGCX \
68+
--sizes=1048576,16777216,104857600 --iters=20
69+
```
70+
71+
### Mooncake
72+
73+
**Server:**
74+
```bash
75+
python kv_transfer_benchmark.py --connector=mooncake --role=server \
76+
--remote-ip=10.8.2.169 --device=gpu \
77+
--sizes=1048576,16777216,104857600 --iters=20
78+
```
79+
80+
**Client:**
81+
```bash
82+
python kv_transfer_benchmark.py --connector=mooncake --role=client \
83+
--remote-ip=10.8.2.169 --device=gpu \
84+
--sizes=1048576,16777216,104857600 --iters=20
85+
```
86+
87+
### FlagCX (Direct)
88+
89+
**Server:**
90+
```bash
91+
cd ~/WORKDIR/FlagCX
92+
python3 test/perf/kv_transfer/kv_transfer_benchmark.py --connector=flagcx --role=server \
93+
--remote-ip=10.8.2.169 --device=gpu \
94+
--sizes=1048576,16777216,104857600 --iters=20
95+
```
96+
97+
**Client:**
98+
```bash
99+
cd ~/WORKDIR/FlagCX
100+
python3 test/perf/kv_transfer/kv_transfer_benchmark.py --connector=flagcx --role=client \
101+
--remote-ip=10.8.2.169 --device=gpu \
102+
--sizes=1048576,16777216,104857600 --iters=20
103+
```
104+
105+
**Custom FlagCX paths:**
106+
```bash
107+
python3 test/perf/kv_transfer/kv_transfer_benchmark.py --connector=flagcx --role=server \
108+
--remote-ip=10.8.2.169 --device=gpu \
109+
--flagcx-lib-path=/custom/path/libflagcx.so \
110+
--flagcx-wrapper-path=/custom/path/FlagCX
111+
```
112+
113+
## Command-Line Arguments
114+
115+
| Argument | Description | Default |
116+
|----------|-------------|---------|
117+
| `--connector` | Transport backend: `nixl`, `mooncake`, or `flagcx` | **required** |
118+
| `--role` | `server` or `client` | **required** |
119+
| `--remote-ip` | Server IP address | `0.0.0.0` |
120+
| `--device` | Memory device: `cpu` or `gpu` | `gpu` |
121+
| `--local-gpu-idx` | CUDA device index | `0` |
122+
| `--sizes` | Comma-separated transfer sizes in bytes | `1KB,4KB,16KB,64KB,256KB,1MB,4MB,16MB,64MB` |
123+
| `--iters` | Number of timed iterations per size | `100` |
124+
| `--warmup-iters` | Number of warmup iterations before timed runs | `10` |
125+
| `--num-blocks` | Number of tensor blocks per transfer | `1` |
126+
| `--nixl-backend` | NIXL backend plugin (e.g., `UCX`, `UCCL`, `FLAGCX`) | `UCX` |
127+
| `--zmq-port` | ZMQ coordination port | `9000` |
128+
| `--mooncake-protocol` | Mooncake protocol: `rdma` or `tcp` | `rdma` |
129+
| `--flagcx-lib-path` | Path to `libflagcx.so` | `$FLAGCX_PATH/build/lib/libflagcx.so` |
130+
| `--flagcx-wrapper-path` | Path to FlagCX root directory | `$FLAGCX_PATH` |
131+
132+
**Note**: Operation type (read/write) is automatically determined by the connector:
133+
- **NIXL**: Uses `read` (client reads from server)
134+
- **Mooncake**: Uses `write` (client writes to server)
135+
- **FlagCX**: Uses `write` (client writes to server)
136+
137+
## Output Format
138+
139+
```
140+
KV Transfer Benchmark connector=nixl role=client
141+
device=gpu gpu=0 op=read iters=100 warmup=10 num_blocks=1
142+
NOTE: nixl connector uses op_type=read (matching vLLM connector semantics)
143+
nixl-backend=UCX
144+
sizes: 1.0 KB, 4.0 KB, 16.0 KB, 64.0 KB, 256.0 KB, 1.0 MB, 4.0 MB, 16.0 MB, 64.0 MB
145+
------------------------------------------------------------------------
146+
1.0 KB | lat= 0.123 ms | BW= 0.01 GB/s ( 0.08 Gbps) | iters=100
147+
✓ Verification passed
148+
4.0 KB | lat= 0.130 ms | BW= 0.03 GB/s ( 0.25 Gbps) | iters=100
149+
✓ Verification passed
150+
16.0 KB | lat= 0.145 ms | BW= 0.11 GB/s ( 0.88 Gbps) | iters=100
151+
✓ Verification passed
152+
64.0 KB | lat= 0.178 ms | BW= 0.36 GB/s ( 2.88 Gbps) | iters=100
153+
✓ Verification passed
154+
256.0 KB | lat= 0.234 ms | BW= 1.09 GB/s ( 8.72 Gbps) | iters=100
155+
✓ Verification passed
156+
1.0 MB | lat= 0.456 ms | BW= 2.19 GB/s ( 17.52 Gbps) | iters=100
157+
✓ Verification passed
158+
4.0 MB | lat= 1.234 ms | BW= 3.24 GB/s ( 25.92 Gbps) | iters=100
159+
✓ Verification passed
160+
16.0 MB | lat= 4.567 ms | BW= 3.50 GB/s ( 28.00 Gbps) | iters=100
161+
✓ Verification passed
162+
64.0 MB | lat= 18.234 ms | BW= 3.51 GB/s ( 28.08 Gbps) | iters=100
163+
✓ Verification passed
164+
------------------------------------------------------------------------
165+
Done.
166+
```
167+
168+
## Verification
169+
170+
The benchmark always performs strict element-by-element verification after each transfer size. Verification runs after timing completes, so it does not affect performance measurements. If any mismatch is detected, an `AssertionError` is raised with detailed diagnostics (which block, which element, expected vs actual value).
171+
172+
Verification semantics per connector:
173+
- **NIXL (read)**: Client checks that received data matches server's pattern (0s)
174+
- **Mooncake (write)**: Server checks that received data matches client's pattern (1s)
175+
- **FlagCX (write)**: Server checks that received data matches client's pattern (1s)
176+
177+
## Architecture
178+
179+
```
180+
┌─────────────────────────────────────────────┐
181+
│ kv_transfer_benchmark.py │
182+
├─────────────────────────────────────────────┤
183+
│ TransportBenchmark (ABC) │
184+
│ setup() → init transport, register mem │
185+
│ run_transfer() → execute one transfer │
186+
│ verify() → check correctness │
187+
│ teardown() → cleanup │
188+
├─────────────────────────────────────────────┤
189+
│ NixlBenchmark │
190+
│ Uses nixl_agent + nixl_agent_config │
191+
│ ZMQ for metadata exchange │
192+
│ Per-iteration handle init/release │
193+
├─────────────────────────────────────────────┤
194+
│ MooncakeBenchmark │
195+
│ Uses mooncake.engine.TransferEngine │
196+
│ ZMQ for session/address exchange │
197+
│ batch_transfer_sync_write() │
198+
├─────────────────────────────────────────────┤
199+
│ FlagCXBenchmark │
200+
│ Uses FLAGCXLibrary (ctypes wrapper) │
201+
│ ZMQ for unique ID exchange │
202+
│ flagcxPut/flagcxBatchPut + signal/wait │
203+
└─────────────────────────────────────────────┘
204+
```
205+
206+
## Comparison with vLLM Connectors
207+
208+
This benchmark measures **transport-level** performance without vLLM overhead:
209+
- No scheduler/worker coordination
210+
- No KV cache block management
211+
- No attention metadata
212+
- Direct memory registration and transfer
213+
214+
For **end-to-end** vLLM KV transfer benchmarks, use:
215+
```bash
216+
vllm serve <model> --kv-transfer-config '{...}'
217+
vllm bench serve --base-url http://127.0.0.1:8000 ...
218+
```
219+
220+
## Troubleshooting
221+
222+
### NIXL Import Error
223+
```
224+
Failed to import NIXL. Is the nixl Python package installed?
225+
```
226+
**Solution**: Install NIXL with `pip install nixl-cu12` (or `nixl-cu11`)
227+
228+
### FLAGCX Backend Not Found
229+
```
230+
NIXL agent initialization failed
231+
```
232+
**Solution**: Ensure `libflagcx.so` is in `LD_LIBRARY_PATH` and `libplugin_FLAGCX.so` is on the NIXL plugin path.
233+
234+
### Mooncake Import Error
235+
```
236+
Failed to import Mooncake.
237+
```
238+
**Solution**: Install Mooncake following https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md
239+
240+
### FlagCX Import Error
241+
```
242+
FLAGCX_PATH not set and --flagcx-wrapper-path not provided.
243+
```
244+
**Solution**: Set `FLAGCX_PATH` environment variable or use `--flagcx-wrapper-path` to point to the FlagCX root directory.
245+
246+
### FlagCX Library Not Found
247+
```
248+
OSError: libflagcx.so: cannot open shared object file
249+
```
250+
**Solution**: Build FlagCX and ensure `libflagcx.so` exists at `$FLAGCX_PATH/build/lib/libflagcx.so`, or use `--flagcx-lib-path` to specify the path.
251+
252+
### ZMQ Connection Timeout
253+
**Solution**: Ensure both server and client use the same `--remote-ip` and `--zmq-port`. Check firewall rules.
254+
255+
### CUDA Out of Memory
256+
**Solution**: Reduce `--sizes` or use `--device=cpu` for testing.
257+
258+
## Plotting Results
259+
260+
Use `plot_benchmark.py` to generate comparison figures from benchmark log files.
261+
262+
### Save benchmark output to log files
263+
264+
```bash
265+
# Run benchmarks and save output
266+
python3 kv_transfer_benchmark.py --connector=nixl --role=server \
267+
--remote-ip=10.8.2.169 --device=gpu --nixl-backend=FLAGCX \
268+
--sizes=1024,4096,16384,65536,262144,1048576,16777216 --iters=20 \
269+
| tee nixl_flagcx.log
270+
271+
python3 kv_transfer_benchmark.py --connector=mooncake --role=server \
272+
--remote-ip=10.8.2.169 --device=gpu \
273+
--sizes=1024,4096,16384,65536,262144,1048576,16777216 --iters=20 \
274+
| tee mooncake.log
275+
276+
python3 kv_transfer_benchmark.py --connector=flagcx --role=server \
277+
--remote-ip=10.8.2.169 --device=gpu \
278+
--sizes=1024,4096,16384,65536,262144,1048576,16777216 --iters=20 \
279+
| tee flagcx.log
280+
```
281+
282+
### Generate comparison plot
283+
284+
```bash
285+
# Auto-discover all *.log files in current directory
286+
python3 plot_benchmark.py
287+
288+
# Specify log directory and output path
289+
python3 plot_benchmark.py --log-dir=./results --output=comparison.png
290+
291+
# Display interactively
292+
python3 plot_benchmark.py --show
293+
```
294+
295+
The script auto-detects connector labels from log file headers (`connector=XXX`, `nixl-backend=YYY`) and produces a side-by-side figure with latency and bandwidth subplots.

0 commit comments

Comments
 (0)