#!/usr/bin/env python3
"""Micro-benchmark for the pre-attention CPU-gap fixes (task #14).

Scenario mirrors the traced production batch: bs=12, prefix 640..26304 tok,
extends ~3.7K tok, cp8 in-seq-split (2*cp segments/request), page 64.

  1. page-aligned validator skip on CUDA tensors
     (memory_pool_host._get_indexer_page_indices hot path) — measured with a
     busy GPU queue, because torch.any/.equal sync for the whole queue.
  2. ragged index descriptor plan: rebuild-per-layer (old) vs per-batch cache.
  3. slot-span builders: rebuild-per-layer (old) vs per-batch cache.

Run (single GPU is enough):
  PYTHONPATH=python python test/manual/bench_cpu_gap_fixes.py
"""
from __future__ import annotations

import time
from types import SimpleNamespace

import torch

from sglang.srt.layers.attention.nsa import nsa_indexer
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
    build_batch_current_slot_spans,
    build_batch_prefix_slot_spans,
    get_or_build_batch_slot_spans,
)
from sglang.srt.mem_cache.page_index_utils import (
    validate_page_aligned_token_indices,
)

DEV = torch.device("cuda", 0)
PAGE = 64
BS = 12
CP = 8
PREFIX_LENS = [19200, 256] + [26304] * 10
EXTEND_LENS = [3776, 7360] + [3347] * 10
REPS = 200


def timed(fn, reps=REPS, warmup=20):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(reps):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / reps * 1e6  # us


def make_fake_batch():
    fb = SimpleNamespace()
    fb.seq_lens_cpu = torch.tensor(
        [p + e for p, e in zip(PREFIX_LENS, EXTEND_LENS)], dtype=torch.int64
    )
    fb.extend_seq_lens_cpu = list(EXTEND_LENS)
    fb.extend_prefix_lens_cpu = list(PREFIX_LENS)
    return fb


def make_cp_index():
    # 2*CP zigzag segments per request over the extend, page-aligned-ish.
    cp_index = []
    for req, extend in enumerate(EXTEND_LENS):
        seg = max(PAGE, (extend // (2 * CP)) // PAGE * PAGE)
        pos = 0
        while pos < extend:
            end = min(pos + seg, extend)
            cp_index.append((req, pos, end))
            pos = end
    return cp_index


def bench_validator():
    n_pages = 600  # ~one layer-group submit worth of pages
    starts = torch.arange(n_pages, device=DEV, dtype=torch.int64) * PAGE
    indices = (
        starts[:, None] + torch.arange(PAGE, device=DEV, dtype=torch.int64)
    ).reshape(-1)

    # Busy queue: enqueue ~0.5ms of GEMM before each validator call, the way
    # the real submit lands behind a layer's compute.
    a = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)
    b = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)

    def old_path():
        for _ in range(4):
            a @ b
        validate_page_aligned_token_indices(indices, PAGE, "bench")
        starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
        return starts2

    def new_path():
        for _ in range(4):
            a @ b
        if not indices.is_cuda:
            validate_page_aligned_token_indices(indices, PAGE, "bench")
        starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
        return starts2

    # measure WALL time per call without trailing torch.cuda.synchronize in
    # the loop (the sync inside the validator is exactly what we measure).
    def wall(fn, reps=60, warmup=10):
        for _ in range(warmup):
            fn()
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(reps):
            fn()
        t1 = time.perf_counter()  # NO sync: CPU-side blocking is the metric
        torch.cuda.synchronize()
        return (t1 - t0) / reps * 1e6

    old = wall(old_path)
    new = wall(new_path)
    print(
        f"1. validator (busy queue, {n_pages} pages): old={old:8.1f}us "
        f"new={new:8.1f}us  speedup x{old/new:.1f}"
    )


def bench_ragged_plan():
    fb = make_fake_batch()
    cp_index = make_cp_index()

    def old_build():
        # the pre-fix per-layer work: full python loop + 7 tensor H2Ds
        return nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)

    fb2 = make_fake_batch()

    def cached():
        return nsa_indexer._get_or_build_cp_ragged_index_plan(
            fb2, cp_index, DEV, None
        )

    old = timed(old_build)
    new = timed(cached)
    n_tokens = sum(EXTEND_LENS)
    print(
        f"2. ragged index plan (bs={BS}, {len(cp_index)} segs, {n_tokens} q tok): "
        f"per-layer rebuild={old:8.1f}us cached={new:8.1f}us  speedup x{old/new:.0f}"
    )


def bench_spans():
    fb = make_fake_batch()
    pages_per_req = max(
        (p + e + PAGE - 1) // PAGE for p, e in zip(PREFIX_LENS, EXTEND_LENS)
    )
    logical_pages = torch.zeros((BS, pages_per_req), dtype=torch.int64)

    def old_build():
        prefix = build_batch_prefix_slot_spans(
            logical_pages=logical_pages,
            prefix_lens_cpu=PREFIX_LENS,
            page_size=PAGE,
        )
        current = build_batch_current_slot_spans(
            logical_pages=logical_pages,
            prefix_lens_cpu=PREFIX_LENS,
            extend_lens_cpu=EXTEND_LENS,
            page_size=PAGE,
        )
        return prefix, current

    def cached():
        return get_or_build_batch_slot_spans(
            fb,
            logical_pages=logical_pages,
            prefix_lens_cpu=PREFIX_LENS,
            extend_lens_cpu=EXTEND_LENS,
            page_size=PAGE,
            want_prefix=True,
        )

    old = timed(old_build, reps=2000)
    new = timed(cached, reps=2000)
    print(
        f"3. slot spans (bs={BS}): per-layer rebuild={old:8.1f}us "
        f"cached={new:8.1f}us  speedup x{old/new:.0f}"
    )


def main():
    torch.cuda.init()
    print(f"device: {torch.cuda.get_device_name(0)}")
    bench_validator()
    bench_ragged_plan()
    bench_spans()
    # equality check: cached plan tensors match a fresh build
    fb = make_fake_batch()
    cp_index = make_cp_index()
    p1 = nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)
    p2 = nsa_indexer._get_or_build_cp_ragged_index_plan(fb, cp_index, DEV, None)
    assert torch.equal(p1.topk_indices_offset_override, p2.topk_indices_offset_override)
    assert torch.equal(p1.kv_lens, p2.kv_lens) and torch.equal(p1.q_bases, p2.q_bases)
    assert p1.segment_records == p2.segment_records
    s1 = build_batch_current_slot_spans(
        logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
        prefix_lens_cpu=PREFIX_LENS,
        extend_lens_cpu=EXTEND_LENS,
        page_size=PAGE,
    )
    _, s2 = get_or_build_batch_slot_spans(
        SimpleNamespace(),
        logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
        prefix_lens_cpu=PREFIX_LENS,
        extend_lens_cpu=EXTEND_LENS,
        page_size=PAGE,
        want_prefix=False,
    )
    assert s1 == s2
    print("EQUALITY CHECKS PASS")


if __name__ == "__main__":
    main()