sgl-project--sglang/python/sglang/jit_kernel/benchmark/kv_canary/utils.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import torch

from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource

BS_AXIS: list[int] = [1, 4, 32, 128, 256, 1024]
PREFIX_AXIS: list[int] = [0, 128, 1024, 4096, 10240, 16384]
EXTEND_LEN_AXIS: list[int] = [128, 512, 4096, 16384]
POOL_AXIS: list[str] = ["full", "swa_window_128"]
REAL_KV_AXIS: list[str] = ["none", "small_1src", "med_2src", "max_4src"]
HASH_MODE_AXIS: list[str] = ["none", "partial", "all"]
SWA_WINDOW: int = 128
RING_CAPACITY: int = 256
MAX_EXTEND_TOKENS_PER_FORWARD: int = 4096


@dataclass(frozen=True, slots=True, kw_only=True)
class BenchCase:
    scenario: str
    bs: int
    prefix_len: int
    mode: str
    extend_len: int
    pool_kind: str
    real_kv_kind: str
    hash_mode: str

    @property
    def case_id(self) -> str:
        return (
            f"{self.scenario}_bs{self.bs}_prefix{self.prefix_len}_{self.mode}{self.extend_len}"
            f"_{self.pool_kind}_rkv{self.real_kv_kind}_hash{self.hash_mode}"
        )


def _case(
    *,
    scenario: str,
    bs: int,
    prefix_len: int,
    mode: str,
    extend_len: int,
    pool_kind: str,
    real_kv_kind: str = "none",
    hash_mode: str = "none",
) -> BenchCase:
    return BenchCase(
        scenario=scenario,
        bs=bs,
        prefix_len=prefix_len,
        mode=mode,
        extend_len=extend_len,
        pool_kind=pool_kind,
        real_kv_kind=real_kv_kind,
        hash_mode=hash_mode,
    )


def _is_realistic_extend_case(case: BenchCase) -> bool:
    if case.mode != "extend":
        return True
    return case.bs * case.extend_len <= MAX_EXTEND_TOKENS_PER_FORWARD


def _dedupe_cases(cases: list[BenchCase]) -> list[BenchCase]:
    seen: set[str] = set()
    result: list[BenchCase] = []

    for case in cases:
        if case.case_id in seen:
            continue
        seen.add(case.case_id)
        result.append(case)

    return result


def build_fast_matrix_cases() -> list[BenchCase]:
    return _dedupe_cases(
        [
            _case(
                scenario="smoke_decode_empty",
                bs=1,
                prefix_len=0,
                mode="decode",
                extend_len=1,
                pool_kind="full",
            ),
            _case(
                scenario="small_extend_batch",
                bs=32,
                prefix_len=4096,
                mode="extend",
                extend_len=128,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_decode_steady",
                bs=256,
                prefix_len=4096,
                mode="decode",
                extend_len=1,
                pool_kind="full",
            ),
            _case(
                scenario="decode_large_batch_short_prefix",
                bs=1024,
                prefix_len=1024,
                mode="decode",
                extend_len=1,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_prefill_chunk_first",
                bs=1,
                prefix_len=0,
                mode="extend",
                extend_len=4096,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_prefill_chunk_mid",
                bs=1,
                prefix_len=8192,
                mode="extend",
                extend_len=4096,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_prefill_chunk_last",
                bs=1,
                prefix_len=12288,
                mode="extend",
                extend_len=4096,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_decode_tail",
                bs=1,
                prefix_len=5120,
                mode="decode",
                extend_len=1,
                pool_kind="full",
            ),
            _case(
                scenario="swa_decode_long_prefix",
                bs=128,
                prefix_len=10240,
                mode="decode",
                extend_len=1,
                pool_kind="swa_window_128",
            ),
            _case(
                scenario="small_extend_single_req",
                bs=1,
                prefix_len=128,
                mode="extend",
                extend_len=128,
                pool_kind="full",
            ),
            _case(
                scenario="medium_extend_chunk",
                bs=4,
                prefix_len=1024,
                mode="extend",
                extend_len=512,
                pool_kind="full",
            ),
            _case(
                scenario="decode_mid_batch",
                bs=128,
                prefix_len=4096,
                mode="decode",
                extend_len=1,
                pool_kind="full",
            ),
            _case(
                scenario="e2e_prefill_chunk_second",
                bs=1,
                prefix_len=4096,
                mode="extend",
                extend_len=4096,
                pool_kind="full",
            ),
            _case(
                scenario="swa_decode_short_prefix",
                bs=256,
                prefix_len=128,
                mode="decode",
                extend_len=1,
                pool_kind="swa_window_128",
            ),
            _case(
                scenario="swa_decode_tail",
                bs=4,
                prefix_len=10240,
                mode="decode",
                extend_len=1,
                pool_kind="swa_window_128",
            ),
            _case(
                scenario="small_extend_batch_hash",
                bs=32,
                prefix_len=4096,
                mode="extend",
                extend_len=128,
                pool_kind="full",
                real_kv_kind="small_1src",
                hash_mode="partial",
            ),
            _case(
                scenario="e2e_prefill_chunk_hash",
                bs=1,
                prefix_len=12288,
                mode="extend",
                extend_len=4096,
                pool_kind="full",
                real_kv_kind="med_2src",
                hash_mode="all",
            ),
            _case(
                scenario="e2e_decode_steady_hash",
                bs=256,
                prefix_len=4096,
                mode="decode",
                extend_len=1,
                pool_kind="full",
                real_kv_kind="max_4src",
                hash_mode="all",
            ),
            _case(
                scenario="swa_decode_long_prefix_hash",
                bs=128,
                prefix_len=10240,
                mode="decode",
                extend_len=1,
                pool_kind="swa_window_128",
                real_kv_kind="med_2src",
                hash_mode="partial",
            ),
            _case(
                scenario="smoke_decode_empty_hash",
                bs=1,
                prefix_len=0,
                mode="decode",
                extend_len=1,
                pool_kind="full",
                real_kv_kind="small_1src",
                hash_mode="all",
            ),
        ]
    )


def build_full_matrix_cases() -> list[BenchCase]:
    """Full matrix plus targeted e2e points.

    Extend cases are pruned to a maximum token chunk per forward because the scheduler chunks long
    prefills; for example, a 4096-token extend is represented as ``bs=1``, not ``bs=32``.
    """
    fast = build_fast_matrix_cases()
    fast_keys = {c.case_id for c in fast}
    full: list[BenchCase] = list(fast)

    for bs in BS_AXIS:
        for prefix_len in PREFIX_AXIS:
            for pool_kind in POOL_AXIS:
                for mode, extend_len in (
                    ("decode", 1),
                    *(("extend", e) for e in EXTEND_LEN_AXIS),
                ):
                    case = _case(
                        scenario="matrix",
                        bs=bs,
                        prefix_len=prefix_len,
                        mode=mode,
                        extend_len=extend_len,
                        pool_kind=pool_kind,
                    )
                    if not _is_realistic_extend_case(case):
                        continue
                    if case.case_id in fast_keys:
                        continue
                    full.append(case)

    fast_base_points = [
        (c.bs, c.prefix_len, c.mode, c.extend_len, c.pool_kind)
        for c in fast
        if c.real_kv_kind == "none" and c.hash_mode == "none"
    ]
    for bs, prefix_len, mode, extend_len, pool_kind in fast_base_points:
        for hash_mode in HASH_MODE_AXIS:
            if hash_mode == "none":
                continue
            for real_kv_kind in REAL_KV_AXIS:
                if real_kv_kind == "none":
                    continue
                case = _case(
                    scenario="fold_matrix",
                    bs=bs,
                    prefix_len=prefix_len,
                    mode=mode,
                    extend_len=extend_len,
                    pool_kind=pool_kind,
                    real_kv_kind=real_kv_kind,
                    hash_mode=hash_mode,
                )
                if not _is_realistic_extend_case(case):
                    continue
                if case.case_id in fast_keys:
                    continue
                full.append(case)
                fast_keys.add(case.case_id)

    return full


def cases_to_x_vals(
    cases: list[BenchCase],
) -> list[tuple[str, int, int, str, int, str, str, str]]:
    return [
        (
            c.scenario,
            c.bs,
            c.prefix_len,
            c.mode,
            c.extend_len,
            c.pool_kind,
            c.real_kv_kind,
            c.hash_mode,
        )
        for c in cases
    ]


def _one_real_kv_source(
    *, num_slots: int, num_bytes: int, read_bytes: int, device: torch.device
) -> RealKvSource:
    tensor = torch.zeros(max(1, num_slots), num_bytes, dtype=torch.uint8, device=device)
    return RealKvSource(
        tensor=tensor,
        page_size=1,
        num_bytes_per_token=num_bytes,
        read_bytes=read_bytes,
    )


def make_real_kv_sources(
    *, kind: str, num_slots: int, device: torch.device
) -> tuple[RealKvSource, ...]:
    """Map a ``real_kv_kind`` axis label to a tuple of ``RealKvSource`` configs.

    Byte-volume ladder (none -> small_1src -> med_2src -> max_4src) so the bench exposes the
    ``real_kv_fold_sources`` PARTIAL/ALL cost gradient. ``max_4src`` hits the
    ``consts.MAX_REAL_KV_SOURCES = 4`` ABI ceiling.
    """
    if kind == "none":
        return ()
    if kind == "small_1src":
        return (
            _one_real_kv_source(
                num_slots=num_slots, num_bytes=16, read_bytes=16, device=device
            ),
        )
    if kind == "med_2src":
        return tuple(
            _one_real_kv_source(
                num_slots=num_slots, num_bytes=32, read_bytes=16, device=device
            )
            for _ in range(2)
        )
    if kind == "max_4src":
        return tuple(
            _one_real_kv_source(
                num_slots=num_slots, num_bytes=64, read_bytes=32, device=device
            )
            for _ in range(4)
        )
    raise ValueError(f"kv-canary bench: unknown real_kv_kind {kind!r}")


def naive_slot_copy_fn(*, total: int, device: torch.device) -> Callable[[], None]:
    n_slots = max(total, 1)
    payload = torch.zeros(n_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device)
    sink = torch.zeros_like(payload)
    indices = torch.arange(n_slots, device=device, dtype=torch.int64) % sink.shape[0]

    def baseline() -> None:
        sink.index_copy_(0, indices, payload)

    return baseline


def naive_cumsum_fn(*, bs: int, device: torch.device) -> Callable[[], None]:
    counts = torch.zeros(max(bs, 1), dtype=torch.int32, device=device)

    def baseline() -> None:
        torch.cumsum(counts, dim=0)

    return baseline