| 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()) |