from __future__ import annotations import csv import platform from dataclasses import dataclass from pathlib import Path import torch import triton.testing from sgl_kernel.kvcacheio import ( transfer_kv_all_layer_lf_pf, transfer_kv_all_layer_mla_lf_pf, ) from sglang.kernels.ops.kvcache.hicache import ( can_use_hicache_jit_kernel, transfer_hicache_all_layer_mla_staged_lf_pf, transfer_hicache_all_layer_staged_lf_pf, ) ROOT = Path(__file__).resolve().parents[2] OUT_DIR = ROOT / "benchmark" / "hicache" / "profiles" CSV_PATH = OUT_DIR / "writeback_compare.csv" MD_PATH = ROOT / "benchmark" / "hicache" / "HICACHE_WRITEBACK_COMPARISON.md" DTYPE = torch.bfloat16 PAGE_SIZE = 64 NUM_LAYERS_LIST = [16, 24, 32, 40, 48, 56, 64, 72, 80] BATCH_PAGES = [4, 8, 16, 32, 64] ELEMENT_DIMS = [256, 512, 1024, 2048] TOTAL_PAGES = max(128, max(BATCH_PAGES) * 2) WARMUP = 5 REP = 25 @dataclass(frozen=True) class BenchRow: mode: str num_layers: int element_dim: int batch_pages: int item_bytes: int baseline_gib_s: float staged_gib_s: float speedup: float def _make_page_starts(total_pages: int, num_pages: int, page_size: int) -> torch.Tensor: order = torch.randperm(total_pages, device="cuda", dtype=torch.int64) return order[:num_pages] * page_size def _make_token_indices( page_starts: torch.Tensor, page_size: int, *, device: str ) -> torch.Tensor: arange = torch.arange(page_size, device=device, dtype=torch.int64) return (page_starts.to(device=device)[:, None] + arange).reshape(-1) def _to_gib_per_s(moved_bytes: int, ms: float) -> float: return moved_bytes / (ms * 1e-3) / (1024**3) def _bench(fn) -> float: ms = triton.testing.do_bench(fn, warmup=WARMUP, rep=REP) return float(ms) def _validate_mha_correctness( baseline, staged, dst_k_baseline: torch.Tensor, dst_v_baseline: torch.Tensor, dst_k_staged: torch.Tensor, dst_v_staged: torch.Tensor, ) -> None: dst_k_baseline.zero_() dst_v_baseline.zero_() dst_k_staged.zero_() dst_v_staged.zero_() baseline() staged() torch.cuda.synchronize() torch.testing.assert_close(dst_k_staged, dst_k_baseline) torch.testing.assert_close(dst_v_staged, dst_v_baseline) def _validate_mla_correctness( baseline, staged, dst_baseline: torch.Tensor, dst_staged: torch.Tensor, ) -> None: dst_baseline.zero_() dst_staged.zero_() baseline() staged() torch.cuda.synchronize() torch.testing.assert_close(dst_staged, dst_baseline) def _bench_mha(num_layers: int, element_dim: int, batch_pages: int) -> BenchRow: item_bytes = element_dim * torch.tensor([], dtype=DTYPE).element_size() assert can_use_hicache_jit_kernel(element_size=item_bytes) total_tokens = TOTAL_PAGES * PAGE_SIZE chunk_tokens = batch_pages * PAGE_SIZE src_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE) dst_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE) src_indices = _make_token_indices(src_page_starts, PAGE_SIZE, device="cuda") dst_indices = _make_token_indices(dst_page_starts, PAGE_SIZE, device="cuda") src_k_layers = [ torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda") for _ in range(num_layers) ] src_v_layers = [ torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda") for _ in range(num_layers) ] src_k_ptrs = torch.tensor( [x.data_ptr() for x in src_k_layers], dtype=torch.uint64, device="cuda" ) src_v_ptrs = torch.tensor( [x.data_ptr() for x in src_v_layers], dtype=torch.uint64, device="cuda" ) dst_k_baseline = torch.empty( total_tokens, num_layers, element_dim, dtype=DTYPE, pin_memory=True ) dst_v_baseline = torch.empty_like(dst_k_baseline, pin_memory=True) dst_k_staged = torch.empty_like(dst_k_baseline, pin_memory=True) dst_v_staged = torch.empty_like(dst_v_baseline, pin_memory=True) staging_k = torch.empty( chunk_tokens, num_layers, element_dim, dtype=DTYPE, device="cuda" ) staging_v = torch.empty_like(staging_k) baseline = lambda: transfer_kv_all_layer_lf_pf( src_k_layers=src_k_ptrs, dst_k=dst_k_baseline, src_v_layers=src_v_ptrs, dst_v=dst_v_baseline, src_indices=src_indices, dst_indices=dst_indices, item_size=item_bytes, dst_layout_dim=item_bytes * num_layers, num_layers=num_layers, ) staged = lambda: transfer_hicache_all_layer_staged_lf_pf( k_ptr_src=src_k_ptrs, v_ptr_src=src_v_ptrs, src_indices=src_indices, dst_indices=dst_indices, staging_k=staging_k, staging_v=staging_v, dst_k=dst_k_staged, dst_v=dst_v_staged, page_size=PAGE_SIZE, ) _validate_mha_correctness( baseline, staged, dst_k_baseline, dst_v_baseline, dst_k_staged, dst_v_staged, ) moved_bytes = chunk_tokens * num_layers * item_bytes * 2 baseline_ms = _bench(baseline) staged_ms = _bench(staged) baseline_bw = _to_gib_per_s(moved_bytes, baseline_ms) staged_bw = _to_gib_per_s(moved_bytes, staged_ms) return BenchRow( mode="MHA", num_layers=num_layers, element_dim=element_dim, batch_pages=batch_pages, item_bytes=item_bytes, baseline_gib_s=baseline_bw, staged_gib_s=staged_bw, speedup=staged_bw / baseline_bw, ) def _bench_mla(num_layers: int, element_dim: int, batch_pages: int) -> BenchRow: item_bytes = element_dim * torch.tensor([], dtype=DTYPE).element_size() assert can_use_hicache_jit_kernel(element_size=item_bytes) total_tokens = TOTAL_PAGES * PAGE_SIZE chunk_tokens = batch_pages * PAGE_SIZE src_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE) dst_page_starts = _make_page_starts(TOTAL_PAGES, batch_pages, PAGE_SIZE) src_indices = _make_token_indices(src_page_starts, PAGE_SIZE, device="cuda") dst_indices = _make_token_indices(dst_page_starts, PAGE_SIZE, device="cuda") src_layers = [ torch.randn(total_tokens, element_dim, dtype=DTYPE, device="cuda") for _ in range(num_layers) ] src_ptrs = torch.tensor( [x.data_ptr() for x in src_layers], dtype=torch.uint64, device="cuda" ) dst_baseline = torch.empty( total_tokens, num_layers, element_dim, dtype=DTYPE, pin_memory=True ) dst_staged = torch.empty_like(dst_baseline, pin_memory=True) staging = torch.empty( chunk_tokens, num_layers, element_dim, dtype=DTYPE, device="cuda" ) baseline = lambda: transfer_kv_all_layer_mla_lf_pf( src_layers=src_ptrs, dst=dst_baseline, src_indices=src_indices, dst_indices=dst_indices, item_size=item_bytes, dst_layout_dim=item_bytes * num_layers, num_layers=num_layers, ) staged = lambda: transfer_hicache_all_layer_mla_staged_lf_pf( ptr_src=src_ptrs, src_indices=src_indices, dst_indices=dst_indices, staging=staging, dst=dst_staged, page_size=PAGE_SIZE, ) _validate_mla_correctness( baseline, staged, dst_baseline, dst_staged, ) moved_bytes = chunk_tokens * num_layers * item_bytes baseline_ms = _bench(baseline) staged_ms = _bench(staged) baseline_bw = _to_gib_per_s(moved_bytes, baseline_ms) staged_bw = _to_gib_per_s(moved_bytes, staged_ms) return BenchRow( mode="MLA", num_layers=num_layers, element_dim=element_dim, batch_pages=batch_pages, item_bytes=item_bytes, baseline_gib_s=baseline_bw, staged_gib_s=staged_bw, speedup=staged_bw / baseline_bw, ) def _write_csv(rows: list[BenchRow]) -> None: OUT_DIR.mkdir(parents=True, exist_ok=True) with CSV_PATH.open("w", newline="") as f: writer = csv.writer(f) writer.writerow( [ "mode", "num_layers", "element_dim", "item_bytes", "batch_pages", "baseline_gib_s", "staged_gib_s", "speedup", ] ) for row in rows: writer.writerow( [ row.mode, row.num_layers, row.element_dim, row.item_bytes, row.batch_pages, f"{row.baseline_gib_s:.6f}", f"{row.staged_gib_s:.6f}", f"{row.speedup:.4f}", ] ) def _table(rows: list[BenchRow], mode: str) -> str: lines = [ "| num_layers | element_dim | item_bytes | batch_pages | `sgl_kernel *_lf_pf` GiB/s | `jit *_staged_lf_pf` GiB/s | speedup |", "| ---: | ---: | ---: | ---: | ---: | ---: | ---: |", ] for row in rows: if row.mode != mode: continue lines.append( f"| {row.num_layers} | {row.element_dim} | {row.item_bytes} | {row.batch_pages} | " f"{row.baseline_gib_s:.3f} | {row.staged_gib_s:.3f} | {row.speedup:.2f}x |" ) return "\n".join(lines) def _summarize(rows: list[BenchRow], mode: str) -> tuple[BenchRow, BenchRow]: filtered = [row for row in rows if row.mode == mode] best = max(filtered, key=lambda row: row.speedup) worst = min(filtered, key=lambda row: row.speedup) return best, worst def _write_report_en(rows: list[BenchRow]) -> None: gpu_name = torch.cuda.get_device_name(0) mha_best, mha_worst = _summarize(rows, "MHA") mla_best, mla_worst = _summarize(rows, "MLA") md = f"""# HiCache Writeback Benchmark Comparison Environment: - GPU: `{gpu_name}` - Platform: `{platform.platform()}` - DType: `{DTYPE}` - `page_size`: `{PAGE_SIZE}` - `num_layers` sweep: `{NUM_LAYERS_LIST}` - `total_pages`: `{TOTAL_PAGES}` - `batch_pages` sweep: `{BATCH_PAGES}` - `element_dim` sweep: `{ELEMENT_DIMS}` - `src_indices` / `dst_indices`: CUDA token indices - Timing: `triton.testing.do_bench(warmup={WARMUP}, rep={REP})` Comparison target: - MHA: `sgl_kernel.transfer_kv_all_layer_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_staged_lf_pf` - MLA: `sgl_kernel.transfer_kv_all_layer_mla_lf_pf` vs `sglang.kernels.ops.kvcache.hicache.transfer_hicache_all_layer_mla_staged_lf_pf` Metric: - Effective writeback bandwidth in `GiB/s` - `speedup = staged_jit / sgl_kernel_lf_pf` ## MHA {_table(rows, "MHA")} ## MLA {_table(rows, "MLA")} ## Takeaways - The staged JIT path is compared directly against the installed `sgl_kernel` LF->PF kernels on the same host-destination writeback workload. - Best MHA staged case: `num_layers={mha_best.num_layers}`, `element_dim={mha_best.element_dim}`, `batch_pages={mha_best.batch_pages}`, `speedup={mha_best.speedup:.2f}x`. - Weakest MHA staged case: `num_layers={mha_worst.num_layers}`, `element_dim={mha_worst.element_dim}`, `batch_pages={mha_worst.batch_pages}`, `speedup={mha_worst.speedup:.2f}x`. - Best MLA staged case: `num_layers={mla_best.num_layers}`, `element_dim={mla_best.element_dim}`, `batch_pages={mla_best.batch_pages}`, `speedup={mla_best.speedup:.2f}x`. - Weakest MLA staged case: `num_layers={mla_worst.num_layers}`, `element_dim={mla_worst.element_dim}`, `batch_pages={mla_worst.batch_pages}`, `speedup={mla_worst.speedup:.2f}x`. - MHA bandwidth uses combined payload bytes for `K + V`. - MLA bandwidth uses single-buffer payload bytes. - Raw data is also available in `benchmark/hicache/profiles/writeback_compare.csv`. """ MD_PATH.write_text(md) def main() -> None: torch.manual_seed(0) rows: list[BenchRow] = [] for num_layers in NUM_LAYERS_LIST: for element_dim in ELEMENT_DIMS: for batch_pages in BATCH_PAGES: rows.append(_bench_mha(num_layers, element_dim, batch_pages)) rows.append(_bench_mla(num_layers, element_dim, batch_pages)) _write_csv(rows) _write_report_en(rows) for row in rows: print( row.mode, "num_layers=", row.num_layers, "element_dim=", row.element_dim, "batch_pages=", row.batch_pages, "baseline_gib_s=", f"{row.baseline_gib_s:.3f}", "staged_gib_s=", f"{row.staged_gib_s:.3f}", "speedup=", f"{row.speedup:.2f}x", ) print(f"wrote {CSV_PATH}") print(f"wrote {MD_PATH}") if __name__ == "__main__": main()