diff --git a/benchmarks/kernels/benchmark_host_buffer_kv_copy.py b/benchmarks/kernels/benchmark_host_buffer_kv_copy.py new file mode 100644 index 000000000000..be4051c0bd41 --- /dev/null +++ b/benchmarks/kernels/benchmark_host_buffer_kv_copy.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the NIXL host-buffer KV copies across KV cache groups. + +`NixlBaseConnectorWorker.sync_recved_kv_to_device` / `save_kv_to_host` move a +request's blocks between the CPU transfer buffer and the device KV cache. Block +ids are unique across KV cache groups, so the per-group copies can be issued as +one. This measures what that coalescing saves: the copied bytes are identical, +only the number of launches changes. +""" + +import time + +import torch +from tabulate import tabulate + +from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks +from vllm.logger import init_logger +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE + +logger = init_logger(__name__) + + +def _make_caches( + num_layers: int, + num_blocks: int, + block_size: int, + num_heads: int, + head_size: int, + dtype: torch.dtype, + device: str, +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Host transfer buffers and device KV caches, shaped as the worker builds them.""" + shape = (num_blocks, block_size, num_heads, head_size) + host = { + f"layer.{i}": torch.randn(shape, dtype=dtype, device="cpu") + for i in range(num_layers) + } + device_caches = { + f"layer.{i}": torch.zeros(shape, dtype=dtype, device=device) + for i in range(num_layers) + } + return host, device_caches + + +def _split_into_groups(block_ids: list[int], num_groups: int) -> list[list[int]]: + """Partition block ids into groups, mirroring the global BlockPool id space.""" + per_group = len(block_ids) // num_groups + return [block_ids[i * per_group : (i + 1) * per_group] for i in range(num_groups)] + + +def _percentile(sorted_values: list[float], q: float) -> float: + return sorted_values[min(len(sorted_values) - 1, int(len(sorted_values) * q))] + + +@torch.inference_mode() +def _time_paired( + host: dict[str, torch.Tensor], + device_caches: dict[str, torch.Tensor], + group_block_ids: list[list[int]], + direction: str, + num_iters: int, +) -> tuple[list[float], list[float]]: + """Per-group and coalesced latencies, measured alternately. + + Interleaving the two variants keeps drift in machine state from landing on + one of them: each iteration times both, back to back. + """ + src, dst = (host, device_caches) if direction == "h2d" else (device_caches, host) + merged = [[b for group in group_block_ids for b in group]] + + def _run(batches: list[list[int]]) -> float: + start = time.perf_counter() + for ids in batches: + copy_kv_blocks(src, dst, ids, ids, direction) + torch.accelerator.synchronize() + return time.perf_counter() - start + + for _ in range(5): # warmup + _run(group_block_ids) + _run(merged) + + per_group: list[float] = [] + coalesced: list[float] = [] + for i in range(num_iters): + # Alternate which variant goes first so ordering cannot favour either. + if i % 2: + per_group.append(_run(group_block_ids)) + coalesced.append(_run(merged)) + else: + coalesced.append(_run(merged)) + per_group.append(_run(group_block_ids)) + return sorted(per_group), sorted(coalesced) + + +@torch.inference_mode() +def _count_device_ops( + host: dict[str, torch.Tensor], + device_caches: dict[str, torch.Tensor], + group_block_ids: list[list[int]], + direction: str, + coalesced: bool, +) -> dict[str, int]: + """Device-op counts for one request-sync, via the profiler. + + Counts do not depend on how busy the machine is, so this is the part of + the comparison that reproduces exactly. + """ + from torch.profiler import ProfilerActivity, profile + + src, dst = (host, device_caches) if direction == "h2d" else (device_caches, host) + batches = ( + [[b for group in group_block_ids for b in group]] + if coalesced + else group_block_ids + ) + + for ids in batches: # warm up allocator/autograd caches + copy_kv_blocks(src, dst, ids, ids, direction) + torch.accelerator.synchronize() + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for ids in batches: + copy_kv_blocks(src, dst, ids, ids, direction) + torch.accelerator.synchronize() + + counts: dict[str, int] = {} + for evt in prof.key_averages(): + if evt.count and (evt.self_device_time_total or evt.device_time_total): + counts[evt.key] = counts.get(evt.key, 0) + evt.count + return counts + + +def _report_op_counts(host, device_caches, block_ids, groups): + print("\nDevice-op counts per request-sync (profiler; independent of load):") + rows = [] + for num_groups in groups: + group_block_ids = _split_into_groups(block_ids, num_groups) + for direction in ("h2d", "d2h"): + before = _count_device_ops( + host, device_caches, group_block_ids, direction, coalesced=False + ) + after = _count_device_ops( + host, device_caches, group_block_ids, direction, coalesced=True + ) + for key in sorted(set(before) | set(after)): + b, a = before.get(key, 0), after.get(key, 0) + if b != a: + rows.append([num_groups, direction, key, b, a]) + print( + tabulate( + rows, + headers=["groups", "direction", "op", "per-group", "coalesced"], + ) + ) + + +def main(args): + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype] + torch.manual_seed(args.seed) + + host, device_caches = _make_caches( + num_layers=args.num_layers, + num_blocks=args.num_blocks, + block_size=args.block_size, + num_heads=args.num_heads, + head_size=args.head_size, + dtype=dtype, + device="cuda", + ) + block_ids = list(range(args.blocks_per_request)) + + rows = [] + for num_groups in args.groups: + group_block_ids = _split_into_groups(block_ids, num_groups) + for direction in ("h2d", "d2h"): + per_group, coalesced = _time_paired( + host, + device_caches, + group_block_ids, + direction, + num_iters=args.iters, + ) + # Interference only ever adds time, so the minimum is the most + # stable estimator here; the p10 spread above it says how noisy + # the machine was while sampling. + best_before = per_group[0] + best_after = coalesced[0] + spread = (_percentile(coalesced, 0.5) - best_after) / best_after + rows.append( + [ + num_groups, + direction, + best_before * 1e6, + best_after * 1e6, + (best_before - best_after) / best_before * 100.0, + spread * 100.0, + ] + ) + + print( + f"layers={args.num_layers} num_blocks={args.num_blocks} " + f"block_size={args.block_size} heads={args.num_heads} " + f"head_size={args.head_size} dtype={args.dtype} " + f"blocks/request={args.blocks_per_request} iters={args.iters}" + ) + print( + tabulate( + rows, + headers=[ + "groups", + "direction", + "per-group min (µs)", + "coalesced min (µs)", + "saved (%)", + "p50 over min (%)", + ], + floatfmt=".3f", + ) + ) + + if args.op_counts: + _report_op_counts(host, device_caches, block_ids, args.groups) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser() + parser.add_argument("--num-layers", type=int, default=48) + parser.add_argument("--num-blocks", type=int, default=2048) + parser.add_argument("--block-size", type=int, choices=[16, 32], default=16) + parser.add_argument("--num-heads", type=int, default=8) + parser.add_argument("--head-size", type=int, default=128) + parser.add_argument("--blocks-per-request", type=int, default=60) + parser.add_argument( + "--groups", + type=int, + nargs="+", + default=[1, 2, 3, 6], + help="KV cache group counts to sweep (1 = non-hybrid model).", + ) + parser.add_argument( + "--dtype", + type=str, + choices=["half", "bfloat16", "float"], + default="bfloat16", + ) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument( + "--op-counts", + action="store_true", + help="Also report profiler op counts, which do not vary with load.", + ) + parser.add_argument("--seed", type=int, default=0) + + main(parser.parse_args()) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 629e8c7ecc3f..cb6c601ba311 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -2396,3 +2396,50 @@ def test_push_write_hybrid_mla_replicates_attention(): assert spec.remote_block_ids == [[7, 8], [3]] assert call.kwargs["local_xfer_side_handle"] == local_handle assert call.kwargs["remote_xfer_side_handle"] == remote_handle + + +def _make_host_buffer_worker(copy_op): + """A worker stripped down to what the host-buffer copy paths touch.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + worker = object.__new__(NixlConnectorWorker) + worker.use_host_buffer = True + worker.copy_blocks = copy_op + return worker + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "group_block_ids,expected_ids", + [ + # Single group (non-hybrid model): one copy, as before. + ([[1, 2, 3]], [1, 2, 3]), + # Hybrid model: three groups collapse into a single copy. + ([[1, 2], [3, 4], [5]], [1, 2, 3, 4, 5]), + # An empty group contributes no ids. + ([[1, 2], []], [1, 2]), + ], +) +def test_sync_recved_kv_issues_one_copy_per_request(group_block_ids, expected_ids): + """h2d copies are issued once per request, not once per KV cache group.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ReqMeta + + calls = [] + worker = _make_host_buffer_worker( + lambda src, dst, src_ids, dst_ids, direction: calls.append( + (src_ids, dst_ids, direction) + ) + ) + worker.host_xfer_buffers = {"layer": None} + worker.device_kv_caches = {"layer": None} + + meta = ReqMeta( + local_block_ids=group_block_ids, + local_physical_block_ids=group_block_ids, + tp_size=1, + ) + worker.sync_recved_kv_to_device("req", meta) + + assert calls == [(expected_ids, expected_ids, "h2d")] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 0cbdd786d8c3..48bb5edc7fcc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2702,17 +2702,18 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): assert self.copy_blocks is not None local_block_ids = meta.local_physical_block_ids - # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups - # The h2d block copies below are intentionally synchronous. + # Every KV cache group allocates from the same BlockPool, so block ids + # are unique across groups and the per-group copies can be issued as one. + block_ids = [block_id for group in local_block_ids for block_id in group] + # The h2d block copy below is intentionally synchronous. with gpu_sync_allowed(): - for group_block_ids in local_block_ids: - self.copy_blocks( - self.host_xfer_buffers, - self.device_kv_caches, - group_block_ids, - group_block_ids, - "h2d", - ) + self.copy_blocks( + self.host_xfer_buffers, + self.device_kv_caches, + block_ids, + block_ids, + "h2d", + ) if logger.isEnabledFor(logging.DEBUG): logger.debug( "synced recved kv of request[%s] to device kv buffer," @@ -2740,14 +2741,18 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): ",".join(map(str, meta.local_physical_block_ids)), ) # blocking - for group_block_ids in meta.local_physical_block_ids: - self.copy_blocks( - self.device_kv_caches, - self.host_xfer_buffers, - group_block_ids, - group_block_ids, - "d2h", - ) + block_ids = [ + block_id + for group in meta.local_physical_block_ids + for block_id in group + ] + self.copy_blocks( + self.device_kv_caches, + self.host_xfer_buffers, + block_ids, + block_ids, + "d2h", + ) @cached_property def _attention_kv_caches(self) -> list[torch.Tensor]: