[KV Connector][NIXL] Coalesce host-buffer KV copies across cache groups (#54483)
Signed-off-by: beenpow <57864041+beenpow@users.noreply.github.com> Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
Chanbin Lim committed
Oct 1, 2026 at 11:10 UTC
47de9d4d049831914858404c48827837c2d915db
3 files changed
+326
-18
benchmarks/kernels/benchmark_host_buffer_kv_copy.py
new
+256
@@ -0,0 +1,256 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Benchmark the NIXL host-buffer KV copies across KV cache groups.
4
+
5
+`NixlBaseConnectorWorker.sync_recved_kv_to_device` / `save_kv_to_host` move a
6
+request's blocks between the CPU transfer buffer and the device KV cache. Block
7
+ids are unique across KV cache groups, so the per-group copies can be issued as
8
+one. This measures what that coalescing saves: the copied bytes are identical,
9
+only the number of launches changes.
10
+"""
11
+
12
+import time
13
+
14
+import torch
15
+from tabulate import tabulate
16
+
17
+from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
18
+from vllm.logger import init_logger
19
+from vllm.utils.argparse_utils import FlexibleArgumentParser
20
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
21
+
22
+logger = init_logger(__name__)
23
+
24
+
25
+def _make_caches(
26
+ num_layers: int,
27
+ num_blocks: int,
28
+ block_size: int,
29
+ num_heads: int,
30
+ head_size: int,
31
+ dtype: torch.dtype,
32
+ device: str,
33
+) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
34
+ """Host transfer buffers and device KV caches, shaped as the worker builds them."""
35
+ shape = (num_blocks, block_size, num_heads, head_size)
36
+ host = {
37
+ f"layer.{i}": torch.randn(shape, dtype=dtype, device="cpu")
38
+ for i in range(num_layers)
39
+ }
40
+ device_caches = {
41
+ f"layer.{i}": torch.zeros(shape, dtype=dtype, device=device)
42
+ for i in range(num_layers)
43
+ }
44
+ return host, device_caches
45
+
46
+
47
+def _split_into_groups(block_ids: list[int], num_groups: int) -> list[list[int]]:
48
+ """Partition block ids into groups, mirroring the global BlockPool id space."""
49
+ per_group = len(block_ids) // num_groups
50
+ return [block_ids[i * per_group : (i + 1) * per_group] for i in range(num_groups)]
51
+
52
+
53
+def _percentile(sorted_values: list[float], q: float) -> float:
54
+ return sorted_values[min(len(sorted_values) - 1, int(len(sorted_values) * q))]
55
+
56
+
57
+@torch.inference_mode()
58
+def _time_paired(
59
+ host: dict[str, torch.Tensor],
60
+ device_caches: dict[str, torch.Tensor],
61
+ group_block_ids: list[list[int]],
62
+ direction: str,
63
+ num_iters: int,
64
+) -> tuple[list[float], list[float]]:
65
+ """Per-group and coalesced latencies, measured alternately.
66
+
67
+ Interleaving the two variants keeps drift in machine state from landing on
68
+ one of them: each iteration times both, back to back.
69
+ """
70
+ src, dst = (host, device_caches) if direction == "h2d" else (device_caches, host)
71
+ merged = [[b for group in group_block_ids for b in group]]
72
+
73
+ def _run(batches: list[list[int]]) -> float:
74
+ start = time.perf_counter()
75
+ for ids in batches:
76
+ copy_kv_blocks(src, dst, ids, ids, direction)
77
+ torch.accelerator.synchronize()
78
+ return time.perf_counter() - start
79
+
80
+ for _ in range(5): # warmup
81
+ _run(group_block_ids)
82
+ _run(merged)
83
+
84
+ per_group: list[float] = []
85
+ coalesced: list[float] = []
86
+ for i in range(num_iters):
87
+ # Alternate which variant goes first so ordering cannot favour either.
88
+ if i % 2:
89
+ per_group.append(_run(group_block_ids))
90
+ coalesced.append(_run(merged))
91
+ else:
92
+ coalesced.append(_run(merged))
93
+ per_group.append(_run(group_block_ids))
94
+ return sorted(per_group), sorted(coalesced)
95
+
96
+
97
+@torch.inference_mode()
98
+def _count_device_ops(
99
+ host: dict[str, torch.Tensor],
100
+ device_caches: dict[str, torch.Tensor],
101
+ group_block_ids: list[list[int]],
102
+ direction: str,
103
+ coalesced: bool,
104
+) -> dict[str, int]:
105
+ """Device-op counts for one request-sync, via the profiler.
106
+
107
+ Counts do not depend on how busy the machine is, so this is the part of
108
+ the comparison that reproduces exactly.
109
+ """
110
+ from torch.profiler import ProfilerActivity, profile
111
+
112
+ src, dst = (host, device_caches) if direction == "h2d" else (device_caches, host)
113
+ batches = (
114
+ [[b for group in group_block_ids for b in group]]
115
+ if coalesced
116
+ else group_block_ids
117
+ )
118
+
119
+ for ids in batches: # warm up allocator/autograd caches
120
+ copy_kv_blocks(src, dst, ids, ids, direction)
121
+ torch.accelerator.synchronize()
122
+
123
+ with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
124
+ for ids in batches:
125
+ copy_kv_blocks(src, dst, ids, ids, direction)
126
+ torch.accelerator.synchronize()
127
+
128
+ counts: dict[str, int] = {}
129
+ for evt in prof.key_averages():
130
+ if evt.count and (evt.self_device_time_total or evt.device_time_total):
131
+ counts[evt.key] = counts.get(evt.key, 0) + evt.count
132
+ return counts
133
+
134
+
135
+def _report_op_counts(host, device_caches, block_ids, groups):
136
+ print("\nDevice-op counts per request-sync (profiler; independent of load):")
137
+ rows = []
138
+ for num_groups in groups:
139
+ group_block_ids = _split_into_groups(block_ids, num_groups)
140
+ for direction in ("h2d", "d2h"):
141
+ before = _count_device_ops(
142
+ host, device_caches, group_block_ids, direction, coalesced=False
143
+ )
144
+ after = _count_device_ops(
145
+ host, device_caches, group_block_ids, direction, coalesced=True
146
+ )
147
+ for key in sorted(set(before) | set(after)):
148
+ b, a = before.get(key, 0), after.get(key, 0)
149
+ if b != a:
150
+ rows.append([num_groups, direction, key, b, a])
151
+ print(
152
+ tabulate(
153
+ rows,
154
+ headers=["groups", "direction", "op", "per-group", "coalesced"],
155
+ )
156
+ )
157
+
158
+
159
+def main(args):
160
+ dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
161
+ torch.manual_seed(args.seed)
162
+
163
+ host, device_caches = _make_caches(
164
+ num_layers=args.num_layers,
165
+ num_blocks=args.num_blocks,
166
+ block_size=args.block_size,
167
+ num_heads=args.num_heads,
168
+ head_size=args.head_size,
169
+ dtype=dtype,
170
+ device="cuda",
171
+ )
172
+ block_ids = list(range(args.blocks_per_request))
173
+
174
+ rows = []
175
+ for num_groups in args.groups:
176
+ group_block_ids = _split_into_groups(block_ids, num_groups)
177
+ for direction in ("h2d", "d2h"):
178
+ per_group, coalesced = _time_paired(
179
+ host,
180
+ device_caches,
181
+ group_block_ids,
182
+ direction,
183
+ num_iters=args.iters,
184
+ )
185
+ # Interference only ever adds time, so the minimum is the most
186
+ # stable estimator here; the p10 spread above it says how noisy
187
+ # the machine was while sampling.
188
+ best_before = per_group[0]
189
+ best_after = coalesced[0]
190
+ spread = (_percentile(coalesced, 0.5) - best_after) / best_after
191
+ rows.append(
192
+ [
193
+ num_groups,
194
+ direction,
195
+ best_before * 1e6,
196
+ best_after * 1e6,
197
+ (best_before - best_after) / best_before * 100.0,
198
+ spread * 100.0,
199
+ ]
200
+ )
201
+
202
+ print(
203
+ f"layers={args.num_layers} num_blocks={args.num_blocks} "
204
+ f"block_size={args.block_size} heads={args.num_heads} "
205
+ f"head_size={args.head_size} dtype={args.dtype} "
206
+ f"blocks/request={args.blocks_per_request} iters={args.iters}"
207
+ )
208
+ print(
209
+ tabulate(
210
+ rows,
211
+ headers=[
212
+ "groups",
213
+ "direction",
214
+ "per-group min (µs)",
215
+ "coalesced min (µs)",
216
+ "saved (%)",
217
+ "p50 over min (%)",
218
+ ],
219
+ floatfmt=".3f",
220
+ )
221
+ )
222
+
223
+ if args.op_counts:
224
+ _report_op_counts(host, device_caches, block_ids, args.groups)
225
+
226
+
227
+if __name__ == "__main__":
228
+ parser = FlexibleArgumentParser()
229
+ parser.add_argument("--num-layers", type=int, default=48)
230
+ parser.add_argument("--num-blocks", type=int, default=2048)
231
+ parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
232
+ parser.add_argument("--num-heads", type=int, default=8)
233
+ parser.add_argument("--head-size", type=int, default=128)
234
+ parser.add_argument("--blocks-per-request", type=int, default=60)
235
+ parser.add_argument(
236
+ "--groups",
237
+ type=int,
238
+ nargs="+",
239
+ default=[1, 2, 3, 6],
240
+ help="KV cache group counts to sweep (1 = non-hybrid model).",
241
+ )
242
+ parser.add_argument(
243
+ "--dtype",
244
+ type=str,
245
+ choices=["half", "bfloat16", "float"],
246
+ default="bfloat16",
247
+ )
248
+ parser.add_argument("--iters", type=int, default=100)
249
+ parser.add_argument(
250
+ "--op-counts",
251
+ action="store_true",
252
+ help="Also report profiler op counts, which do not vary with load.",
253
+ )
254
+ parser.add_argument("--seed", type=int, default=0)
255
+
256
+ main(parser.parse_args())
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+47
@@ -2360,3 +2360,50 @@ def test_push_write_hybrid_mla_replicates_attention():
2360
assert spec.remote_block_ids == [[7, 8], [3]]
2361
assert call.kwargs["local_xfer_side_handle"] == local_handle
2362
assert call.kwargs["remote_xfer_side_handle"] == remote_handle
2363
+
2364
+
2365
+def _make_host_buffer_worker(copy_op):
2366
+ """A worker stripped down to what the host-buffer copy paths touch."""
2367
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
2368
+ NixlConnectorWorker,
2369
+ )
2370
+
2371
+ worker = object.__new__(NixlConnectorWorker)
2372
+ worker.use_host_buffer = True
2373
+ worker.copy_blocks = copy_op
2374
+ return worker
2375
+
2376
+
2377
+@pytest.mark.cpu_test
2378
+@pytest.mark.parametrize(
2379
+ "group_block_ids,expected_ids",
2380
+ [
2381
+ # Single group (non-hybrid model): one copy, as before.
2382
+ ([[1, 2, 3]], [1, 2, 3]),
2383
+ # Hybrid model: three groups collapse into a single copy.
2384
+ ([[1, 2], [3, 4], [5]], [1, 2, 3, 4, 5]),
2385
+ # An empty group contributes no ids.
2386
+ ([[1, 2], []], [1, 2]),
2387
+ ],
2388
+)
2389
+def test_sync_recved_kv_issues_one_copy_per_request(group_block_ids, expected_ids):
2390
+ """h2d copies are issued once per request, not once per KV cache group."""
2391
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ReqMeta
2392
+
2393
+ calls = []
2394
+ worker = _make_host_buffer_worker(
2395
+ lambda src, dst, src_ids, dst_ids, direction: calls.append(
2396
+ (src_ids, dst_ids, direction)
2397
+ )
2398
+ )
2399
+ worker.host_xfer_buffers = {"layer": None}
2400
+ worker.device_kv_caches = {"layer": None}
2401
+
2402
+ meta = ReqMeta(
2403
+ local_block_ids=group_block_ids,
2404
+ local_physical_block_ids=group_block_ids,
2405
+ tp_size=1,
2406
+ )
2407
+ worker.sync_recved_kv_to_device("req", meta)
2408
+
2409
+ assert calls == [(expected_ids, expected_ids, "h2d")]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+23
-18
@@ -2702,17 +2702,18 @@ class NixlBaseConnectorWorker:
2702
assert self.copy_blocks is not None
2703
2704
local_block_ids = meta.local_physical_block_ids
2705
- # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups
2706
- # The h2d block copies below are intentionally synchronous.
2705
+ # Every KV cache group allocates from the same BlockPool, so block ids
2706
+ # are unique across groups and the per-group copies can be issued as one.
2707
+ block_ids = [block_id for group in local_block_ids for block_id in group]
2708
+ # The h2d block copy below is intentionally synchronous.
2709
with gpu_sync_allowed():
2708
- for group_block_ids in local_block_ids:
2709
- self.copy_blocks(
2710
- self.host_xfer_buffers,
2711
- self.device_kv_caches,
2712
- group_block_ids,
2713
- group_block_ids,
2714
- "h2d",
2715
- )
2710
+ self.copy_blocks(
2711
+ self.host_xfer_buffers,
2712
+ self.device_kv_caches,
2713
+ block_ids,
2714
+ block_ids,
2715
+ "h2d",
2716
+ )
2717
if logger.isEnabledFor(logging.DEBUG):
2718
logger.debug(
2719
"synced recved kv of request[%s] to device kv buffer,"
@@ -2740,14 +2741,18 @@ class NixlBaseConnectorWorker:
2741
",".join(map(str, meta.local_physical_block_ids)),
2742
)
2743
# blocking
2743
- for group_block_ids in meta.local_physical_block_ids:
2744
- self.copy_blocks(
2745
- self.device_kv_caches,
2746
- self.host_xfer_buffers,
2747
- group_block_ids,
2748
- group_block_ids,
2749
- "d2h",
2750
- )
2744
+ block_ids = [
2745
+ block_id
2746
+ for group in meta.local_physical_block_ids
2747
+ for block_id in group
2748
+ ]
2749
+ self.copy_blocks(
2750
+ self.device_kv_caches,
2751
+ self.host_xfer_buffers,
2752
+ block_ids,
2753
+ block_ids,
2754
+ "d2h",
2755
+ )
2756
2757
@cached_property
2758
def _attention_kv_caches(self) -> list[torch.Tensor]: