diff --git a/benchmark/hicache/bench_cp_hicache_backup_overlap.py b/benchmark/hicache/bench_cp_hicache_backup_overlap.py new file mode 100644 index 000000000..b5c26a800 --- /dev/null +++ b/benchmark/hicache/bench_cp_hicache_backup_overlap.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +"""Compare short-extend CP HiCache backup strategies. + +The per-layer path models current direct+page_first_direct forward-overlap +backup. The sync fallback baseline uses a TAI SM-kernel all-layer LF->PF +path for the same page_first_direct host layout, because a synchronous tail +copy can spend SMs to reduce latency and improve bandwidth. +""" + +import argparse +import gc +import statistics +import time +from dataclasses import dataclass +from typing import Callable + +import torch + +from tai_kernel.nsa_prefill import ( + transfer_kv_all_layer_page_first_direct_lf_pf_from_ptrs, + transfer_kv_per_layer_direct_lf_pf, +) + + +@dataclass +class RunResult: + total_ms: float + forward_ms: float + extra_ms: float + tail_ms: float + submit_ms: float + + +def _parse_int_list(value: str) -> list[int]: + return [int(item) for item in value.split(",") if item] + + +def _parse_float_list(value: str) -> list[float]: + return [float(item) for item in value.split(",") if item] + + +def _median_result(results: list[RunResult]) -> RunResult: + return RunResult( + total_ms=statistics.median(r.total_ms for r in results), + forward_ms=statistics.median(r.forward_ms for r in results), + extra_ms=statistics.median(r.extra_ms for r in results), + tail_ms=statistics.median(r.tail_ms for r in results), + submit_ms=statistics.median(r.submit_ms for r in results), + ) + + +def _make_page_indices( + pool_tokens: int, transfer_tokens: int, page_size: int, *, random_pages: bool +) -> torch.Tensor: + if transfer_tokens % page_size != 0: + raise ValueError("transfer_tokens must be divisible by page_size") + if pool_tokens % page_size != 0: + raise ValueError("pool_tokens must be divisible by page_size") + transfer_pages = transfer_tokens // page_size + total_pages = pool_tokens // page_size + if transfer_pages > total_pages: + raise ValueError("transfer_tokens must be <= pool_tokens") + if random_pages: + page_ids = torch.randperm(total_pages, dtype=torch.long)[:transfer_pages] + else: + page_ids = torch.arange(transfer_pages, dtype=torch.long) + return torch.cat( + [ + torch.arange( + int(page_id) * page_size, + (int(page_id) + 1) * page_size, + dtype=torch.long, + ) + for page_id in page_ids + ] + ).contiguous() + + +def _event_elapsed_ms(start: torch.cuda.Event, end: torch.cuda.Event) -> float: + return float(start.elapsed_time(end)) + + +def _sleep_on_stream(cycles: int) -> None: + if cycles > 0: + torch.cuda._sleep(cycles) + + +def _measure_sleep_ms(cycles: int, repeat: int = 5) -> float: + samples: list[float] = [] + stream = torch.cuda.current_stream() + for _ in range(repeat): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + _sleep_on_stream(cycles) + end.record(stream) + end.synchronize() + samples.append(_event_elapsed_ms(start, end)) + return statistics.median(samples) + + +def _calibrate_sleep_cycles(target_us: float) -> int: + if target_us <= 0: + return 0 + target_ms = target_us / 1000.0 + cycles = 100_000 + measured = _measure_sleep_ms(cycles, repeat=3) + while measured < max(target_ms * 0.25, 0.01): + cycles *= 4 + measured = _measure_sleep_ms(cycles, repeat=3) + for _ in range(4): + if measured <= 0: + break + cycles = max(1, int(cycles * target_ms / measured)) + measured = _measure_sleep_ms(cycles, repeat=3) + return cycles + + +def _build_pools( + *, + layers: int, + pool_tokens: int, + page_size: int, + item_size: int, + num_tensors: int, + dtype: torch.dtype, +) -> tuple[ + list[list[torch.Tensor]], + torch.Tensor, + list[torch.Tensor], +]: + total_pages = pool_tokens // page_size + src_by_tensor: list[list[torch.Tensor]] = [] + direct_host_buffers: list[torch.Tensor] = [] + for _ in range(num_tensors): + storage = torch.empty( + (layers, pool_tokens, item_size), dtype=dtype, device="cuda" + ).contiguous() + storage.normal_() + src_layers = [storage[layer] for layer in range(layers)] + src_by_tensor.append(src_layers) + # direct+page_first_direct layout: [page, layer, token-in-page, item]. + direct_host_buffers.append( + torch.empty( + (total_pages, layers, page_size, item_size), + dtype=dtype, + device="cpu", + pin_memory=True, + ) + ) + src_layer_ptrs = torch.tensor( + [src.data_ptr() for tensor_layers in src_by_tensor for src in tensor_layers], + dtype=torch.uint64, + device="cuda", + ) + return src_by_tensor, src_layer_ptrs, direct_host_buffers + + +def _run_forward_only(layers: int, sleep_cycles: int) -> RunResult: + stream = torch.cuda.current_stream() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record(stream) + for _ in range(layers): + _sleep_on_stream(sleep_cycles) + end.record(stream) + end.synchronize() + total = _event_elapsed_ms(start, end) + return RunResult( + total_ms=total, + forward_ms=total, + extra_ms=0.0, + tail_ms=0.0, + submit_ms=0.0, + ) + + +def _run_all_layer_sync( + *, + src_layer_ptrs: torch.Tensor, + host_buffers: list[torch.Tensor], + src_indices_cuda: torch.Tensor, + dst_indices_cuda: torch.Tensor, + layers: int, + page_size: int, + item_size_bytes: int, + sleep_cycles: int, +) -> RunResult: + compute_stream = torch.cuda.current_stream() + copy_stream = torch.cuda.Stream() + start = torch.cuda.Event(enable_timing=True) + forward_done = torch.cuda.Event(enable_timing=True) + copy_done = torch.cuda.Event(enable_timing=True) + + start.record(compute_stream) + for _ in range(layers): + _sleep_on_stream(sleep_cycles) + forward_done.record(compute_stream) + + submit_start = time.perf_counter() + with torch.cuda.stream(copy_stream): + copy_stream.wait_event(forward_done) + transfer_kv_all_layer_page_first_direct_lf_pf_from_ptrs( + src_layer_ptrs, + host_buffers, + src_indices_cuda, + dst_indices_cuda, + page_size=page_size, + item_size=item_size_bytes, + num_layers=layers, + ) + copy_done.record(copy_stream) + submit_ms = (time.perf_counter() - submit_start) * 1000.0 + copy_done.synchronize() + + forward_ms = _event_elapsed_ms(start, forward_done) + total_ms = _event_elapsed_ms(start, copy_done) + extra_ms = max(0.0, total_ms - forward_ms) + return RunResult( + total_ms=total_ms, + forward_ms=forward_ms, + extra_ms=extra_ms, + tail_ms=extra_ms, + submit_ms=submit_ms, + ) + + +def _run_per_layer_async( + *, + src_by_tensor: list[list[torch.Tensor]], + host_buffers: list[torch.Tensor], + src_indices_cpu: torch.Tensor, + dst_indices_cpu: torch.Tensor, + layers: int, + page_size: int, + sleep_cycles: int, +) -> RunResult: + compute_stream = torch.cuda.current_stream() + copy_stream = torch.cuda.Stream() + start = torch.cuda.Event(enable_timing=True) + forward_done = torch.cuda.Event(enable_timing=True) + copy_done = torch.cuda.Event(enable_timing=True) + submit_ms = 0.0 + + start.record(compute_stream) + for layer_id in range(layers): + _sleep_on_stream(sleep_cycles) + layer_done = torch.cuda.Event() + layer_done.record(compute_stream) + + submit_start = time.perf_counter() + with torch.cuda.stream(copy_stream): + copy_stream.wait_event(layer_done) + transfer_kv_per_layer_direct_lf_pf( + [tensor_layers[layer_id] for tensor_layers in src_by_tensor], + host_buffers, + src_indices_cpu, + dst_indices_cpu, + layer_id=layer_id, + page_size=page_size, + ) + submit_ms += (time.perf_counter() - submit_start) * 1000.0 + + forward_done.record(compute_stream) + with torch.cuda.stream(copy_stream): + copy_done.record(copy_stream) + forward_done.synchronize() + copy_done.synchronize() + + forward_ms = _event_elapsed_ms(start, forward_done) + total_ms = _event_elapsed_ms(start, copy_done) + tail_ms = max(0.0, total_ms - forward_ms) + return RunResult( + total_ms=total_ms, + forward_ms=forward_ms, + extra_ms=tail_ms, + tail_ms=tail_ms, + submit_ms=submit_ms, + ) + + +def _time_mode( + fn: Callable[[], RunResult], *, warmup: int, repeat: int +) -> RunResult: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + results = [fn() for _ in range(repeat)] + torch.cuda.synchronize() + return _median_result(results) + + +def _benchmark_case( + *, + tokens: int, + forward_us: float, + args: argparse.Namespace, +) -> None: + pool_tokens = args.pool_tokens or tokens + dtype = getattr(torch, args.dtype) + item_size_bytes = args.item_size * torch.empty((), dtype=dtype).element_size() + sleep_cycles = _calibrate_sleep_cycles(forward_us) + src_indices_cpu = _make_page_indices( + pool_tokens, tokens, args.page_size, random_pages=args.random_pages + ) + dst_indices_cpu = _make_page_indices( + pool_tokens, tokens, args.page_size, random_pages=args.random_pages + ) + src_indices_cuda = src_indices_cpu.to("cuda") + dst_indices_cuda = dst_indices_cpu.to("cuda") + src_by_tensor, src_layer_ptrs, direct_host_buffers = _build_pools( + layers=args.layers, + pool_tokens=pool_tokens, + page_size=args.page_size, + item_size=args.item_size, + num_tensors=args.num_tensors, + dtype=dtype, + ) + + forward = _time_mode( + lambda: _run_forward_only(args.layers, sleep_cycles), + warmup=args.warmup, + repeat=args.repeat, + ) + all_layer = _time_mode( + lambda: _run_all_layer_sync( + src_layer_ptrs=src_layer_ptrs, + host_buffers=direct_host_buffers, + src_indices_cuda=src_indices_cuda, + dst_indices_cuda=dst_indices_cuda, + layers=args.layers, + page_size=args.page_size, + item_size_bytes=item_size_bytes, + sleep_cycles=sleep_cycles, + ), + warmup=args.warmup, + repeat=args.repeat, + ) + per_layer = _time_mode( + lambda: _run_per_layer_async( + src_by_tensor=src_by_tensor, + host_buffers=direct_host_buffers, + src_indices_cpu=src_indices_cpu, + dst_indices_cpu=dst_indices_cpu, + layers=args.layers, + page_size=args.page_size, + sleep_cycles=sleep_cycles, + ), + warmup=args.warmup, + repeat=args.repeat, + ) + all_over_forward_only_ms = max(0.0, all_layer.total_ms - forward.total_ms) + per_over_forward_only_ms = max(0.0, per_layer.total_ms - forward.total_ms) + winner = "per_layer_async" if per_layer.total_ms < all_layer.total_ms else "all_layer_sync" + bytes_per_layer = tokens * item_size_bytes * args.num_tensors + total_backup_bytes = bytes_per_layer * args.layers + print( + "RESULT " + f"tokens={tokens} forward_us_per_layer={forward_us:.3f} " + f"sleep_cycles={sleep_cycles} layers={args.layers} " + f"num_tensors={args.num_tensors} item_size={args.item_size} " + f"item_size_bytes={item_size_bytes} dtype={args.dtype} " + f"random_pages={args.random_pages} " + f"all_layer_backend=tai_kernel_page_first_direct " + f"per_layer_backend=direct_page_first_direct " + f"backup_total_mb={total_backup_bytes / 1e6:.2f} " + f"forward_only_ms={forward.total_ms:.3f} " + f"all_total_ms={all_layer.total_ms:.3f} " + f"all_extra_ms={all_layer.extra_ms:.3f} " + f"all_over_forward_only_ms={all_over_forward_only_ms:.3f} " + f"all_submit_ms={all_layer.submit_ms:.3f} " + f"per_total_ms={per_layer.total_ms:.3f} " + f"per_extra_ms={per_layer.extra_ms:.3f} " + f"per_over_forward_only_ms={per_over_forward_only_ms:.3f} " + f"per_tail_ms={per_layer.tail_ms:.3f} " + f"per_submit_ms={per_layer.submit_ms:.3f} " + f"delta_total_ms={per_layer.total_ms - all_layer.total_ms:.3f} " + f"delta_over_forward_only_ms={per_over_forward_only_ms - all_over_forward_only_ms:.3f} " + f"winner={winner}" + ) + + del ( + src_by_tensor, + src_layer_ptrs, + direct_host_buffers, + src_indices_cpu, + dst_indices_cpu, + src_indices_cuda, + dst_indices_cuda, + ) + gc.collect() + torch.cuda.empty_cache() + + +def main() -> None: + parser = argparse.ArgumentParser( + "Benchmark CP HiCache backup strategy: per-layer async " + "direct+page_first_direct forward-overlap vs sync all-layer " + "TAI kernel+page_first_direct tail backup." + ) + parser.add_argument("--tokens", default="1024,4096,10240") + parser.add_argument("--forward-us-per-layer", default="0,50,100,200") + parser.add_argument("--layers", type=int, default=78) + parser.add_argument("--pool-tokens", type=int, default=0) + parser.add_argument("--page-size", type=int, default=64) + parser.add_argument("--item-size", type=int, default=576) + parser.add_argument("--num-tensors", type=int, choices=[1, 2], default=1) + parser.add_argument( + "--dtype", choices=["float16", "bfloat16", "float32"], default="bfloat16" + ) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeat", type=int, default=5) + parser.add_argument("--random-pages", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if not hasattr(torch.cuda, "_sleep"): + raise RuntimeError("torch.cuda._sleep is required for synthetic forward") + + torch.manual_seed(20260528) + for tokens in _parse_int_list(args.tokens): + for forward_us in _parse_float_list(args.forward_us_per_layer): + _benchmark_case(tokens=tokens, forward_us=forward_us, args=args) + + +if __name__ == "__main__": + main()