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