lightseekorg--tokenspeed/test/runtime/benchmark/bench_decode_cache.py

"""
Benchmark the cache operations that run on the hot decode path.

During decode, for each completed page tokenspeed calls:
  insert(paged_key, page_ids)
  match_prefix(paged_key)             # to confirm hit
  dec_lock_ref(old_last_node)
  inc_lock_ref(new_last_node)

We simulate a steady-state server: a large cache of pre-existing sequences
(simulating the populated prefix tree after many requests) plus N concurrent
"requests" each decoding D pages one at a time.

Run with:
  python/.venv/bin/python test/runtime/benchmark/bench_decode_cache.py
"""

from __future__ import annotations

import dataclasses
import time
from typing import Optional

import torch

from tokenspeed.runtime.cache.prefix_cache import CacheInitParams, PrefixCache, TreeNode

# ---------------------------------------------------------------------------
# Minimal mock allocator — matches what real decode needs
# ---------------------------------------------------------------------------


class MockAllocator:
    def __init__(self, page_size: int = 16):
        self.page_size = page_size
        self._next_page = 0
        # Lazily allocated on first access; avoids 128MB upfront per benchmark run.
        self._req_to_page: torch.Tensor | None = None

    @property
    def req_to_page(self) -> torch.Tensor:
        if self._req_to_page is None:
            self._req_to_page = torch.zeros(512, 512, dtype=torch.int32)
        return self._req_to_page

    @property
    def req_to_page_cpu(self) -> torch.Tensor:
        return self.req_to_page

    def alloc_page(self) -> int:
        p = self._next_page
        self._next_page += 1
        return p

    # Called by insert/match_prefix indirectly
    def free_with_diff(self, new_ids, old_ids):
        diff = new_ids != old_ids
        return diff

    def append_to_later_free(self, t):
        pass

    def free_group_end(self):
        pass


def _make_cache(page_size: int = 16) -> tuple[PrefixCache, MockAllocator]:
    alloc = MockAllocator(page_size=page_size)
    params = CacheInitParams(
        disable=False,
        req_to_token_pool=None,
        token_to_kv_pool_allocator=alloc,
        page_size=page_size,
    )
    return PrefixCache(params), alloc


# ---------------------------------------------------------------------------
# Scenario builders
# ---------------------------------------------------------------------------


def _populate_cache(
    cache: PrefixCache, alloc: MockAllocator, n_background: int, prefix_pages: int = 0
) -> None:
    """Fill cache with n_background independent sequences (background traffic)."""
    page_size = cache.page_size
    for s in range(n_background):
        if prefix_pages > 0:
            key = [tuple([0] * page_size)] * prefix_pages + [
                tuple([s + 1] + [0] * (page_size - 1))
            ]
        else:
            key = [tuple([s + 1] + [0] * (page_size - 1))]
        n_pages = len(key)
        value = torch.tensor(
            [alloc.alloc_page() for _ in range(n_pages)], dtype=torch.int32
        )
        cache.insert(key, value)


# ---------------------------------------------------------------------------
# Decode simulation
# ---------------------------------------------------------------------------


@dataclasses.dataclass
class SimReq:
    req_id: int
    prefix_key: list  # shared prefix already in cache
    decode_pages: int  # how many pages this req will generate
    current_pages: int = 0
    last_node: Optional[TreeNode] = None


def _simulate_decode_step(
    cache: PrefixCache, alloc: MockAllocator, req: SimReq
) -> bool:
    """
    Simulate one page completion for req.
    Returns True if the request is done.
    """
    page_size = cache.page_size
    req.current_pages += 1

    # Build key: shared prefix + unique decode pages so far
    decode_part = [
        tuple([req.req_id * 10000 + p + 1] + [0] * (page_size - 1))
        for p in range(req.current_pages)
    ]
    full_key = req.prefix_key + decode_part

    value = torch.tensor(
        [alloc.alloc_page() for _ in range(len(full_key))], dtype=torch.int32
    )

    # Hot path: insert → match → dec_lock_ref → inc_lock_ref
    cache.insert(full_key, value)
    result = cache.match_prefix(full_key)
    new_last_node = result.last_device_node

    cache.dec_lock_ref(req.last_node)
    cache.inc_lock_ref(new_last_node)
    req.last_node = new_last_node

    return req.current_pages >= req.decode_pages


# ---------------------------------------------------------------------------
# Benchmark runner
# ---------------------------------------------------------------------------


def bench(
    label: str,
    n_background: int,
    n_concurrent: int,
    decode_pages: int,
    prefix_pages: int = 0,
    page_size: int = 16,
    repeats: int = 5,
) -> dict:
    """
    Measure total time for n_concurrent requests each decoding decode_pages pages,
    against a background cache of n_background entries.
    """
    times = []
    for _ in range(repeats):
        cache, alloc = _make_cache(page_size)
        _populate_cache(cache, alloc, n_background, prefix_pages)

        # Build shared prefix key (same for all requests)
        if prefix_pages > 0:
            prefix_key = [tuple([0] * page_size)] * prefix_pages
        else:
            prefix_key = []

        requests = [
            SimReq(req_id=i, prefix_key=prefix_key, decode_pages=decode_pages)
            for i in range(n_concurrent)
        ]

        # Lock each request's last_node at their prefix (simulating prefill done)
        for req in requests:
            if prefix_key:
                result = cache.match_prefix(prefix_key)
                req.last_node = result.last_device_node
                cache.inc_lock_ref(req.last_node)

        total_steps = n_concurrent * decode_pages
        step = 0
        active = list(requests)

        t0 = time.perf_counter()
        while active:
            next_active = []
            for req in active:
                done = _simulate_decode_step(cache, alloc, req)
                step += 1
                if not done:
                    next_active.append(req)
            active = next_active
        elapsed = time.perf_counter() - t0

        times.append(elapsed)

    times.sort()
    median = times[len(times) // 2]
    total_steps = n_concurrent * decode_pages
    return {
        "label": label,
        "n_background": n_background,
        "n_concurrent": n_concurrent,
        "decode_pages": decode_pages,
        "prefix_pages": prefix_pages,
        "total_decode_steps": total_steps,
        "median_ms": round(median * 1e3, 3),
        "us_per_step": round(median * 1e6 / total_steps, 2),
    }


if __name__ == "__main__":
    import json

    scenarios = [
        # (label,                   n_background, n_concurrent, decode_pages, prefix_pages)
        ("small_cache_short_seq", 1_000, 64, 32, 0),
        ("small_cache_long_seq", 1_000, 64, 128, 0),
        ("large_cache_short_seq", 20_000, 64, 32, 0),
        ("large_cache_long_seq", 20_000, 64, 128, 0),
        ("large_cache_shared_pfx", 20_000, 64, 32, 100),
        ("xlarge_cache_short_seq", 50_000, 64, 32, 0),
    ]

    for args in scenarios:
        r = bench(*args)
        print(json.dumps(r))