Ground HiCache backup overlap policy in measured tails

The per-layer backup path did not explain end-to-end performance by inspection alone, so add a focused benchmark that compares all-layer tail backup against per-layer forward-overlap backup under the same page_first_direct layout. The benchmark reports total latency, over-forward-only latency, submit overhead, and winner so short-extend fallback decisions are based on measured cost rather than request-size intuition. Constraint: CUDA benchmark execution is remote-only; this local commit only records the benchmark driver. Rejected: Adding more online logs for this question | the user asked for benchmark evidence and logs would perturb scheduler behavior. Confidence: high Scope-risk: narrow Directive: Do not infer production policy from per_extra alone; per-layer CPU submission can be hidden inside the forward timing interval, so total latency is the safer comparison metric. Tested: Local py_compile benchmark/hicache/bench_cp_hicache_backup_overlap.py. Tested: Remote benchmark logs generated under /mnt/beegfs/cjy/cp_hicache_backup_overlap_bench_20260527_*.log. Not-tested: Full SGLang test suite; benchmark is standalone and CUDA execution was remote-only.
2026-05-28 03:43:05 +08:00
parent 67d52346de
commit 2c94b8de23
1 changed files with 429 additions and 0 deletions
--- a/benchmark/hicache/bench_cp_hicache_backup_overlap.py
+++ 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()