Blame for src/tinyllm/engine.py

"""Step 4.4: serving many requests: static batching vs continuous batching.

Workload model: offline, every request is queued at t=0 (no arrival process yet).

- run_static: take requests in order, B at a time. Left-pad the prompts, prefill once,
  decode until the batch's *longest* request is done. Rows that finish early keep
  computing (wasted work), and the next batch waits for the slowest row.
- run_continuous: S slots (SlotKVCache). Every step, queued requests are admitted into
  free slots (one batched, padded prefill for the newcomers), then every slot decodes
  one token together. A finished request frees its slot at once.

Both return per-request Results (generated tokens, TTFT, latency) and Stats (throughput,
wasted decode rows, KV memory reserved vs used). Correctness: each request's tokens must
equal its solo generate() output (tests/test_engine.py).
"""

import math
import time
from collections import deque
from dataclasses import dataclass, field

import torch
from torch import Tensor

from tinyllm.generate import sample_next
from tinyllm.kv_cache import KVCache, SlotKVCache
from tinyllm.model import Transformer


@dataclass(frozen=True)
class Sampling:
    """How to pick tokens (sample_next's settings). Frozen, so equal settings group together."""
    temperature: float = 0.0             # 0 = greedy
    top_k: int | None = None
    top_p: float | None = None

    def __post_init__(self):
        if not self.temperature >= 0:
            raise ValueError(f"temperature must be >= 0, got {self.temperature}")
        if self.top_k is not None and self.top_k < 1:
            raise ValueError(f"top_k must be >= 1, got {self.top_k}")
        if self.top_p is not None and not 0.0 < self.top_p <= 1.0:
            raise ValueError(f"top_p must be in (0, 1], got {self.top_p}")


@dataclass
class Request:
    prompt: list[int]
    max_new_tokens: int
    eos_id: int | None = None
    sampling: Sampling | None = None     # ContinuousEngine only; None = the engine's settings


@dataclass
class Result:
    tokens: list[int]          # generated tokens only (EOS included if it was produced)
    ttft_s: float              # engine start → first generated token
    latency_s: float           # engine start → last generated token


@dataclass
class Stats:
    wall_s: float = 0.0
    useful_tokens: int = 0     # generated tokens that belong to a request
    decode_steps: int = 0
    computed_rows: int = 0     # rows run through decode (incl. finished rows / free slots)
    useful_rows: int = 0       # rows whose output was kept
    kv_reserved_tokens: int = 0            # KV capacity held (max over the run)
    kv_used_frac: list[float] = field(default_factory=list)  # per decode step: real tokens / capacity
    kv_alloc_frac: list[float] = field(default_factory=list)  # per decode step: real tokens / claimed
    kv_bytes_per_token: int = 0
    preemptions: int = 0
    peak_active: int = 0
    kv_allocated_at_end: int = 0           # leak check: must be 0 once every request is done

    def summary(self, results: list[Result]) -> dict:
        lat = sorted(r.latency_s for r in results)
        ttft = sorted(r.ttft_s for r in results)
        p95 = lambda xs: xs[min(len(xs) - 1, math.ceil(0.95 * len(xs)) - 1)]
        return {
            "requests": len(results),
            "wall_s": self.wall_s,
            "useful_tokens": self.useful_tokens,
            "tok_s": self.useful_tokens / self.wall_s,
            "decode_steps": self.decode_steps,
            "row_util": self.useful_rows / max(1, self.computed_rows),
            "latency_mean_s": sum(lat) / len(lat),
            "latency_p95_s": p95(lat),
            "ttft_mean_s": sum(ttft) / len(ttft),
            "ttft_p95_s": p95(ttft),
            "kv_reserved_mib": self.kv_reserved_tokens * self.kv_bytes_per_token / 2**20,
            "kv_used_mean": sum(self.kv_used_frac) / max(1, len(self.kv_used_frac)),
            "kv_used_peak": max(self.kv_used_frac, default=0.0),
            "kv_alloc_used_mean": (sum(self.kv_alloc_frac) / len(self.kv_alloc_frac)
                                   if self.kv_alloc_frac else None),
            "preemptions": self.preemptions,
            "peak_active": self.peak_active,
            "kv_allocated_at_end": self.kv_allocated_at_end,
        }


def _autocast(device: torch.device):
    return torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda")


def _kv_bytes_per_token(model: Transformer, device: torch.device) -> int:
    elem = 2 if device.type == "cuda" else 4          # bf16 under autocast, else fp32
    return sum(layer.self_attn.cache_bytes_per_token(elem) for layer in model.layers)


@torch.no_grad()
def prefill_padded(model: Transformer, prompts: list[list[int]], max_len: int,
                   pad_id: int = 0) -> tuple[Tensor, KVCache, list[int]]:
    """Left-pad prompts to one length, run one prefill. Returns last-position logits (B, V),
    the filled cache (pos = padded length) and each row's pad count."""
    device = next(model.parameters()).device
    L = max(len(p) for p in prompts)
    pad = [L - len(p) for p in prompts]
    idx = torch.tensor([[pad_id] * n + p for n, p in zip(pad, prompts)], device=device)
    cache = KVCache(model.cfg.n_layers, max_len, pad_lens=torch.tensor(pad, device=device))
    with _autocast(device):
        logits, _ = model(idx, cache=cache)
    return logits[:, -1], cache, pad


def _done(req: Request, toks: list[int]) -> bool:
    return len(toks) >= req.max_new_tokens or (req.eos_id is not None and toks[-1] == req.eos_id)


@torch.no_grad()
def run_static(model: Transformer, requests: list[Request], batch_size: int,
               temperature: float = 0.0, top_k=None, top_p=None, generator=None,
               ) -> tuple[list[Result], dict]:
    if any(r.sampling is not None for r in requests):
        raise ValueError("run_static uses one sampling setting; per-request sampling needs ContinuousEngine")
    model.eval()
    device = next(model.parameters()).device
    pick = dict(temperature=temperature, top_k=top_k, top_p=top_p, generator=generator)
    stats = Stats(kv_bytes_per_token=_kv_bytes_per_token(model, device))
    results: list[Result | None] = [None] * len(requests)
    t0 = time.perf_counter()
    for b0 in range(0, len(requests), batch_size):
        batch = requests[b0:b0 + batch_size]
        L = max(len(r.prompt) for r in batch)
        new_max = max(r.max_new_tokens for r in batch)
        if L + new_max > model.cfg.max_seq_len:
            raise ValueError(f"padded prompt {L} + new {new_max} exceeds max_seq_len")
        logits, cache, _ = prefill_padded(model, [r.prompt for r in batch], L + new_max)
        capacity = len(batch) * (L + new_max)
        stats.kv_reserved_tokens = max(stats.kv_reserved_tokens, capacity)
        outs: list[list[int]] = [[] for _ in batch]
        ttft = [0.0] * len(batch)
        done = [False] * len(batch)
        while True:
            nxt = sample_next(logits, **pick)
            toks = nxt.squeeze(-1).tolist()                      # GPU sync: tokens to the host
            now = time.perf_counter() - t0
            for j, (r, t) in enumerate(zip(batch, toks)):
                if done[j]:
                    continue
                outs[j].append(t)
                if len(outs[j]) == 1:
                    ttft[j] = now
                if _done(r, outs[j]):
                    done[j] = True
                    results[b0 + j] = Result(outs[j], ttft[j], now)
            if all(done):
                break
            live = [j for j in range(len(batch)) if not done[j]]
            stats.decode_steps += 1
            stats.computed_rows += len(batch)
            stats.useful_rows += len(live)
            stats.kv_used_frac.append(
                sum(len(batch[j].prompt) + len(outs[j]) for j in live) / capacity)
            with _autocast(device):
                logits, _ = model(nxt, cache=cache)
            logits = logits[:, -1]
    stats.wall_s = time.perf_counter() - t0
    stats.useful_tokens = sum(len(r.tokens) for r in results)
    return results, stats.summary(results)


@dataclass
class Event:
    """One generated token, as step() reports it (a server streams these)."""
    rid: int
    token: int
    done: bool                 # the request finished with this token (EOS or max_new_tokens)


class ContinuousEngine:
    """Step 4.13: run_continuous's loop as an object, so requests can arrive while others
    decode. add_request() queues a request at any time; each step() admits queued requests
    into free slots (one padded prefill), decodes one token for every active slot and
    returns the new tokens as Events. A finished request's Result is in `results` (pop it).

    Times are seconds on the engine clock (since construction). TTFT and latency count from
    the request's arrival, so with every request added at once they match 4.4's numbers.
    Sampling: the engine's settings are the default; a request's `sampling` overrides them.
    Rows with equal settings are sampled in one call, so when every row uses the same
    settings the result (and RNG use) is exactly as with per-engine settings. One generator
    is shared by all rows. Not thread-safe: call add_request / step / abort from one thread.

    kv="slots": SlotKVCache, every slot reserves max_len positions (4.4).
    kv="paged": PagedKVCache with kv_blocks blocks of block_size tokens (4.5). Default
    kv_blocks = the same memory as the slot cache. Sequences beyond what the blocks hold are
    preempted and recomputed later.
    cuda_graph (4.11b, CUDA, kv="slots" only): the decode step is replayed as a CUDA graph per
    attention-window bucket (cuda_graph.SlotGraphDecoder); prefill stays eager. True builds a
    decoder for this engine; pass a SlotGraphDecoder to reuse its cache and captured graphs.
    """

    def __init__(self, model: Transformer, n_slots: int, max_len: int | None = None,
                 temperature: float = 0.0, top_k=None, top_p=None, generator=None,
                 kv: str = "slots", block_size: int = 16, kv_blocks: int | None = None,
                 cuda_graph=False, buckets: tuple[int, ...] | None = None):
        model.eval()
        cfg = model.cfg
        self.model = model
        self.device = device = next(model.parameters()).device
        self.n_slots = n_slots
        self.max_len = max_len = max_len or cfg.max_seq_len
        self.sampling = Sampling(temperature, top_k, top_p)    # the default for every request
        self.generator = generator
        dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        if cuda_graph is not False and kv != "slots":
            raise ValueError("cuda_graph supports kv='slots' only (paged: not yet)")
        self.decoder = None
        if cuda_graph is not False and not isinstance(cuda_graph, bool):   # a SlotGraphDecoder
            self.decoder, slots = cuda_graph, cuda_graph.slots
            if slots.n_slots != n_slots or slots.max_len != max_len or any(slots.active):
                raise ValueError(f"decoder cache has {slots.n_slots} slots of {slots.max_len} "
                                 f"(active: {sum(slots.active)}); need {n_slots} idle slots of {max_len}")
        elif kv == "slots":
            slots = SlotKVCache(cfg.n_layers, n_slots, cfg.n_kv_heads, cfg.head_dim, max_len, dtype, device)
        elif kv == "paged":
            from tinyllm.paged_kv import PagedKVCache
            n_blocks = kv_blocks or n_slots * math.ceil(max_len / block_size)
            slots = PagedKVCache(cfg.n_layers, n_slots, cfg.n_kv_heads, cfg.head_dim, n_blocks,
                                 block_size, max_len, dtype, device)
        else:
            raise ValueError(f"kv must be 'slots' or 'paged', got {kv!r}")
        if cuda_graph is True:
            from tinyllm.cuda_graph import SlotGraphDecoder
            self.decoder = SlotGraphDecoder(model, slots, buckets=buckets)
        self.slots = slots
        self.stats = Stats(kv_bytes_per_token=_kv_bytes_per_token(model, device),
                           kv_reserved_tokens=slots.reserved_tokens())
        self.results: dict[int, Result] = {}              # finished, until the caller pops them
        self._req: dict[int, Request] = {}                # unfinished requests only
        self._arrival: dict[int, float] = {}
        self._outs: dict[int, list[int]] = {}
        self._ttft: dict[int, float] = {}
        self._slot_req: list[int | None] = [None] * n_slots
        self._last_tok = torch.zeros(n_slots, 1, dtype=torch.long, device=device)
        self._queue: deque[int] = deque()
        self._next_id = 0
        self._t0 = time.perf_counter()

    # -- public API ------------------------------------------------------------
    def now(self) -> float:
        return time.perf_counter() - self._t0

    def add_request(self, req: Request) -> int:
        """Queue a request; returns its id. Raises at once if it can never fit a slot."""
        if not req.prompt or req.max_new_tokens < 1:
            raise ValueError("a request needs a non-empty prompt and max_new_tokens >= 1")
        if len(req.prompt) + req.max_new_tokens > self.max_len:
            raise ValueError(f"prompt {len(req.prompt)} + new {req.max_new_tokens} "
                             f"exceeds slot max_len {self.max_len}")
        rid, self._next_id = self._next_id, self._next_id + 1
        self._req[rid] = req
        self._arrival[rid] = self.now()
        self._queue.append(rid)
        return rid

    def has_work(self) -> bool:
        return bool(self._queue) or any(self.slots.active)

    def load(self) -> tuple[int, int]:
        """(requests waiting for a slot, slots decoding a request)."""
        return len(self._queue), sum(self.slots.active)

    def kv_used_frac(self) -> float:
        return self.slots.used_tokens() / max(1, self.slots.reserved_tokens())

    def abort(self, rid: int) -> bool:
        """Drop an unfinished request (e.g. its client went away), freeing its slot.
        Returns False if it is unknown or already finished. No Result is kept."""
        if rid not in self._req:
            return False
        if rid in self._queue:
            self._queue.remove(rid)
        if rid in self._slot_req:
            s = self._slot_req.index(rid)
            self._slot_req[s] = None
            self.slots.release(s)
        self._forget(rid)
        return True

    @torch.no_grad()
    def step(self) -> list[Event]:
        """Admit what fits, decode one token per active slot. Returns the new tokens."""
        events: list[Event] = []
        slots = self.slots
        free = slots.free_slots()
        rids: list[int] = []
        if free and self._queue:                              # admit newcomers: one padded prefill
            q = self._queue
            if not self._fits_padded([q[0]]):
                r = q[0]
                raise ValueError(f"request {r}: prompt {len(self._context(r))} + new {self._left(r)} "
                                 f"exceeds slot max_len {self.max_len}")
            while (q and len(rids) < len(free) and self._fits_padded(rids + [q[0]])
                   and slots.can_admit([len(self._context(r)) for r in rids + [q[0]]])):
                rids.append(q.popleft())
            if not rids and not any(slots.active):
                raise ValueError(f"request {q[0]} does not fit in the KV cache even alone")
        if rids:
            prompts = [self._context(r) for r in rids]
            L = max(len(p) for p in prompts)
            logits, pre, pad = prefill_padded(self.model, prompts, max_len=L)
            nxt = self._sample(logits, [self._sampling(r) for r in rids])
            taken = free[:len(rids)]
            slots.admit(taken, pre, pad)
            self._last_tok[torch.tensor(taken, device=self.device)] = nxt
            toks = nxt.squeeze(-1).tolist()
            now = self.now()
            for s, rid, t in zip(taken, rids, toks):
                self._slot_req[s] = rid
                self._outs.setdefault(rid, [])
                events.append(self._record(s, rid, t, now))
        for s in slots.reserve_step():                        # paged: out of blocks → preempt newest
            self._queue.appendleft(self._slot_req[s])         # recompute later from its context
            self._slot_req[s] = None
            self.stats.preemptions += 1
        if not any(slots.active):
            return events
        st = self.stats
        live = [s for s in range(self.n_slots) if slots.active[s]]
        st.peak_active = max(st.peak_active, len(live))
        st.decode_steps += 1
        st.computed_rows += self.n_slots
        st.useful_rows += len(live)
        used = slots.used_tokens()
        st.kv_used_frac.append(used / slots.reserved_tokens())
        st.kv_alloc_frac.append(used / max(1, slots.allocated_tokens()))
        if self.decoder is not None:
            logits = self.decoder.step(self._last_tok)        # static buffer: sample at once
        else:
            with _autocast(self.device):
                logits, _ = self.model(self._last_tok, cache=slots)   # one token per slot
        rows = [self._sampling(r) if r is not None else self.sampling for r in self._slot_req]
        self._last_tok = self._sample(logits[:, -1], rows)    # idle rows: engine default
        toks = self._last_tok.squeeze(-1).tolist()
        now = self.now()
        for s in live:
            events.append(self._record(s, self._slot_req[s], toks[s], now))
        return events

    # -- internals -------------------------------------------------------------
    def _context(self, rid: int) -> list[int]:
        """Prompt plus tokens generated so far (non-empty only after a preemption)."""
        return self._req[rid].prompt + self._outs.get(rid, [])

    def _sampling(self, rid: int) -> Sampling:
        return self._req[rid].sampling or self.sampling

    def _sample(self, logits: Tensor, rows: list[Sampling]) -> Tensor:
        """Next token per row (B, V) → (B, 1), each row with its own settings. One
        sample_next call per distinct setting, in order of first appearance."""
        groups: dict[Sampling, list[int]] = {}
        for i, s in enumerate(rows):
            groups.setdefault(s, []).append(i)
        if len(groups) == 1:
            (s,) = groups
            return sample_next(logits, s.temperature, s.top_k, s.top_p, self.generator)
        out = torch.empty(logits.size(0), 1, dtype=torch.long, device=logits.device)
        for s, idx in groups.items():
            ix = torch.tensor(idx, device=logits.device)
            out[ix] = sample_next(logits[ix], s.temperature, s.top_k, s.top_p, self.generator)
        return out

    def _left(self, rid: int) -> int:
        return self._req[rid].max_new_tokens - len(self._outs.get(rid, []))

    def _fits_padded(self, rids: list[int]) -> bool:
        """Every row of a padded prefill sits at positions up to L + its remaining tokens, so
        that must fit max_len (the RoPE table, and slot storage). A preempted request comes back
        with a longer context, so this can fail for a batch even when each request fits alone."""
        return max(len(self._context(r)) for r in rids) + max(self._left(r) for r in rids) <= self.max_len

    def _record(self, s: int, rid: int, tok: int, now: float) -> Event:
        outs = self._outs[rid]
        outs.append(tok)
        if len(outs) == 1:
            self._ttft[rid] = now
        done = _done(self._req[rid], outs)
        if done:
            t_in = self._arrival[rid]
            self.results[rid] = Result(outs, self._ttft[rid] - t_in, now - t_in)
            self.stats.useful_tokens += len(outs)
            self._slot_req[s] = None
            self.slots.release(s)
            self._forget(rid)
        return Event(rid, tok, done)

    def _forget(self, rid: int) -> None:
        for d in (self._req, self._arrival, self._outs, self._ttft):
            d.pop(rid, None)


@torch.no_grad()
def run_continuous(model: Transformer, requests: list[Request], n_slots: int,
                   max_len: int | None = None, temperature: float = 0.0, top_k=None,
                   top_p=None, generator=None, kv: str = "slots", block_size: int = 16,
                   kv_blocks: int | None = None, cuda_graph=False,
                   buckets: tuple[int, ...] | None = None) -> tuple[list[Result], dict]:
    """Offline continuous batching (4.4): every request queued at t=0, run to completion.
    A thin wrapper over ContinuousEngine; see it for kv / cuda_graph."""
    eng = ContinuousEngine(model, n_slots, max_len, temperature, top_k, top_p, generator,
                           kv, block_size, kv_blocks, cuda_graph, buckets)
    t0 = time.perf_counter()
    rids = [eng.add_request(r) for r in requests]
    while eng.has_work():
        eng.step()
    st = eng.stats
    st.kv_allocated_at_end = eng.slots.allocated_tokens()
    st.wall_s = time.perf_counter() - t0
    results = [eng.results.pop(r) for r in rids]
    return results, st.summary(results)