chore: import upstream snapshot with attribution

python/sglang/jit_kernel/benchmark/kv_canary/utils.py

@@ -0,0 +1,404 @@
+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