tiny-llm 75e4579: 6.5b step 5: 5a Space created, 5b secrets added

.dockerignore

@@ -0,0 +1,7 @@
+# Allow-list: the image needs only the package and the tokenizer (see Dockerfile).
+# Never data, checkpoints, runs, wandb, .venv, notes or secrets.
+*
+!src/tinyllm/
+!tokenizer/tinystories-bpe8192.json
+**/__pycache__
+**/*.pyc

Dockerfile

@@ -0,0 +1,58 @@
+# The tinyllm server on CPU (6.5b, D-033): the image behind the Hugging Face Space, and a local `docker run`.
+#
+#   docker build -t tinyllm --build-arg GIT_COMMIT=$(git rev-parse --short HEAD) .
+#   docker run --rm -p 7860:7860 tinyllm
+#   curl localhost:7860/v1/completions -H 'content-type: application/json' \
+#        -d '{"prompt": "Once upon a time", "max_tokens": 60}'
+#
+# CPU only: no CUDA, Triton or CUDA graphs here, so the model runs the PyTorch code paths
+# (/stats reports norm_backend "torch"). The speed numbers in the README come from the laptop GPU.
+# No weights in the image: the model is downloaded from the Hub at start into the HF cache
+# (MODEL / REVISION below), or mounted as a folder: -v /path/to/export:/model -e MODEL=/model.
+
+FROM python:3.12-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 PIP_DISABLE_PIP_VERSION_CHECK=1
+
+# CPU build of the torch version in pyproject.toml (as in CI), then only what the server imports.
+# Versions match uv.lock; training-only packages (datasets, wandb) stay out of the image.
+RUN pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cpu \
+ && pip install numpy==2.5.4 tokenizers==0.23.3 huggingface-hub==2.4.0 \
+                fastapi==0.162.0 starlette==1.10.0 pydantic==2.14.2 uvicorn==0.54.0 \
+ && python -c "import torch; assert not torch.cuda.is_available() and torch.version.cuda is None"
+
+# Spaces run the container as uid 1000; the HF cache must be writable by it.
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user HF_HOME=/home/user/.cache/huggingface PYTHONPATH=/app/src
+WORKDIR /app
+
+COPY --chown=user src/tinyllm ./src/tinyllm
+COPY --chown=user tokenizer ./tokenizer
+
+ARG GIT_COMMIT=75e4579
+# REVISION pins the Hub commit of the model upload (scripts/upload_model.py), so a rebuild serves the same weights.
+ENV TINYLLM_GIT_COMMIT=$GIT_COMMIT \
+    MODEL=CoolDip00/tiny-llm-tinystories-30m REVISION=b11100b189f9d5085ce8c0fd775ca71a6eaf287c \
+    SLOTS=8 PORT=7860 SLO_TTFT_S=2 LOG_DATASET="" \
+    MAX_QUEUE=4 RATE_LIMIT=20 RATE_BURST=5 PROXY_HOPS=1
+
+EXPOSE 7860
+# Shell form so every ENV above can be overridden with -e (or a Space variable); exec keeps uvicorn as
+# PID 1 (SIGTERM → graceful drain, then the last log upload). --revision only for a repo id: a mounted
+# folder has none. SLO_TTFT_S: the status page's TTFT SLO threshold, 2 s on the shared CPU.
+# MAX_QUEUE / RATE_LIMIT (requests/min per visitor, 0 = off) / RATE_BURST: load shedding (docs/SERVING.md).
+# SLOTS / MAX_QUEUE / RATE_*: from the 2-CPU load test (docs/SERVING.md "Space defaults"): 8 slots,
+# queue ≈ 2 s TTFT target × throughput, 20 requests/min per visitor with a burst of 5.
+# --trust-proxy only on a Space (HF sets SPACE_ID): there HF's proxy adds the visitor's address to
+# X-Forwarded-For; a plain `docker run` has no proxy, so the header is the client's own and is ignored.
+# PROXY_HOPS: proxies in front that append to that header (the client is that many entries from the right).
+# LOG_DATASET (off when empty): a private HF dataset for the request logs (docs/SERVING.md "Privacy");
+# needs HF_TOKEN with write access to it, and TINYLLM_LOG_SALT to count visitors across restarts.
+CMD exec python -m tinyllm.server --hf-model "$MODEL" \
+    $( [ -d "$MODEL" ] || echo --revision "$REVISION" ) \
+    $( [ -z "$LOG_DATASET" ] || echo --log-dir /tmp/tinyllm-logs --log-dataset "$LOG_DATASET" ) \
+    $( [ -z "$SPACE_ID" ] || echo --trust-proxy --proxy-hops "$PROXY_HOPS" ) \
+    --device cpu --host 0.0.0.0 --port "$PORT" --slots "$SLOTS" --slo-ttft-s "$SLO_TTFT_S" \
+    --max-queue "$MAX_QUEUE" --rate-limit "$RATE_LIMIT" --rate-burst "$RATE_BURST"

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 cooldip00
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR THE OPTIONS IN THE
+SOFTWARE.

README.md

@@ -1,10 +1,30 @@
 ---
-title: Tiny Llm
-emoji: 💻
-colorFrom: blue
-colorTo: yellow
+title: tiny-llm
+emoji: 📖
+colorFrom: indigo
+colorTo: gray
 sdk: docker
+app_port: 7860
+license: mit
+short_description: A 30M-parameter story model, served by a from-scratch engine
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# tiny-llm
+
+A 30M-parameter language model trained from scratch on TinyStories, served by the project's own inference server (continuous batching, streaming, load shedding) on HF's free 2-vCPU CPU tier.
+
+- **Demo:** this page. Type the start of a story and it carries on.
+- **Status:** `/status` shows live latency, throughput and the SLOs (time to first token p95 ≤ 2 s, 99% availability).
+- **API:** `POST /v1/completions` (OpenAI-style), e.g.
+  `curl https://cooldip00-tiny-llm.hf.space/v1/completions -H 'content-type: application/json' -d '{"prompt": "Once upon a time", "max_tokens": 60}'`
+- **Code, design notes and benchmarks:** https://github.com/cooldip00/tiny-llm (the Dockerfile here is the one in the repo; the exact commit is shown on `/stats`).
+- **Model:** https://huggingface.co/CoolDip00/tiny-llm-tinystories-30m
+
+## Limits
+
+It is a small model on a shared CPU: stories are simple and sometimes odd. Each visitor may send 20 requests per minute (bursts of 5); when the server is full it answers 429 at once instead of making everyone wait. The Space sleeps after 48 hours without visitors; the first visit after that takes a while to start.
+
+## What is logged
+
+Request metadata only: timings, token counts, status and finish reason. **Never the prompt or the generated text.** Visitors are counted by a salted hash of the IP address; the IP itself is never stored. The logs go to a private dataset and are kept for 30 days. Details: [docs/SERVING.md](https://github.com/cooldip00/tiny-llm/blob/main/docs/SERVING.md#privacy-rules).

src/tinyllm/__init__.py

@@ -0,0 +1 @@
+"""tinyllm: a small Llama-style LLM built from scratch."""

src/tinyllm/attention/__init__.py

@@ -0,0 +1,7 @@
+"""Attention variants. Importing this package registers every built-in variant."""
+
+from tinyllm.attention.base import ATTENTION, Attention, build_attention, register
+from tinyllm.attention import mha  # noqa: F401  (registers "mha")
+from tinyllm.attention import gqa  # noqa: F401  (registers "gqa")
+
+__all__ = ["ATTENTION", "Attention", "build_attention", "register"]

src/tinyllm/attention/base.py

@@ -0,0 +1,48 @@
+"""The swappable attention contract and the registry (docs/ARCHITECTURE.md).
+
+Every variant (mha, gqa, mla, gdn) subclasses `Attention`, registers itself under
+a name, and is picked per layer through `ModelConfig.layer_types`. The rest of
+the model never needs to know which one it got.
+"""
+
+from torch import Tensor, nn
+
+from tinyllm.config import ModelConfig
+
+ATTENTION: dict[str, type["Attention"]] = {}
+
+
+def register(name: str):
+    """Class decorator: @register("mha") adds the class to ATTENTION."""
+    def deco(cls: type["Attention"]) -> type["Attention"]:
+        if name in ATTENTION:
+            raise ValueError(f"attention variant {name!r} already registered")
+        ATTENTION[name] = cls
+        return cls
+    return deco
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: ModelConfig, layer_idx: int):
+        super().__init__()
+        self.cfg = cfg
+        self.layer_idx = layer_idx
+
+    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
+        """x: (B, T, d_model) -> (B, T, d_model).
+
+        rope: (cos, sin) for the T positions in x.
+        cache: per-layer state for incremental decoding. None during training.
+        """
+        raise NotImplementedError
+
+    def cache_bytes_per_token(self, dtype_bytes: int = 2) -> int:
+        """Decode-time state this layer adds per generated token (bf16 = 2 bytes)."""
+        raise NotImplementedError
+
+
+def build_attention(cfg: ModelConfig, layer_idx: int) -> Attention:
+    name = cfg.layer_types[layer_idx]
+    if name not in ATTENTION:
+        raise KeyError(f"unknown attention {name!r}; registered: {sorted(ATTENTION)}")
+    return ATTENTION[name](cfg, layer_idx)

src/tinyllm/attention/gqa.py

@@ -0,0 +1,65 @@
+"""Grouped-query attention (Ainslie et al. 2023), step 4.9a. TinyLlama: 32 query heads, 4 KV heads.
+
+Query heads are split into n_kv_heads groups of g = n_heads / n_kv_heads; every head in group j
+attends with the same key/value head j. So q_proj stays d -> n_heads·hd while k_proj / v_proj
+shrink to d -> n_kv_heads·hd, and the KV cache stores n_kv_heads heads: g× less decode state.
+
+Head order matches HF Llama's `repeat_kv` (query head i uses KV head i // g), which is what
+TinyLlama's weights expect (4.9b / 4.9c).
+
+With n_kv_heads == n_heads this is exactly `mha` (same parameter names and shapes, so an `mha`
+state dict loads into it).
+"""
+
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from tinyllm.attention.base import Attention, register
+from tinyllm.config import ModelConfig
+from tinyllm.model import apply_rope
+
+
+def repeat_kv(x: Tensor, groups: int) -> Tensor:
+    """(B, kv, T, hd) -> (B, kv·groups, T, hd): KV head j is copied to query heads j·g .. j·g+g-1.
+
+    The explicit form of HF's grouping. Used as the reference in tests; forward uses enable_gqa.
+    """
+    if groups == 1:
+        return x
+    B, kv, T, hd = x.shape
+    return x[:, :, None].expand(B, kv, groups, T, hd).reshape(B, kv * groups, T, hd)
+
+
+@register("gqa")
+class GQA(Attention):
+    def __init__(self, cfg: ModelConfig, layer_idx: int):
+        super().__init__(cfg, layer_idx)
+        d, h, kv, hd = cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.head_dim
+        self.n_heads, self.n_kv_heads, self.head_dim = h, kv, hd
+        self.groups = h // kv
+        self.q_proj = nn.Linear(d, h * hd, bias=False)
+        self.k_proj = nn.Linear(d, kv * hd, bias=False)
+        self.v_proj = nn.Linear(d, kv * hd, bias=False)
+        self.o_proj = nn.Linear(h * hd, d, bias=False)
+
+    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
+        """cache: a kv_cache.LayerKV / SlotLayerKV / paged layer view. It stores n_kv_heads heads."""
+        B, T, _ = x.shape
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        cos, sin = rope
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        mask = {"is_causal": True}
+        if cache is not None:
+            k, v = cache.update(k, v)                    # (B, kv, S, hd): the cache holds only kv heads
+            mask = cache.attn_kwargs
+        # enable_gqa: the kernel maps query head i to KV head i // g itself, without copying K/V.
+        # Measured in 4.9a (torch 2.14, TinyLlama heads, bf16): 0.07-0.74x the time of
+        # repeat_kv + SDPA, masks included. repeat_kv at decode copies the whole cache every
+        # step (B=32, 512 positions: 1.09 ms vs 0.08 ms). repeat_kv stays as the test reference.
+        y = F.scaled_dot_product_attention(q, k, v, enable_gqa=self.groups > 1, **mask)
+        return self.o_proj(y.transpose(1, 2).reshape(B, T, -1))
+
+    def cache_bytes_per_token(self, dtype_bytes: int = 2) -> int:
+        return 2 * self.n_kv_heads * self.head_dim * dtype_bytes  # one K and one V vector per KV head

src/tinyllm/attention/mha.py

@@ -0,0 +1,45 @@
+"""Baseline multi-head causal self-attention (D-010: PyTorch SDPA)."""
+
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from tinyllm.attention.base import Attention, register
+from tinyllm.config import ModelConfig
+from tinyllm.model import apply_rope
+
+
+@register("mha")
+class MHA(Attention):
+    """softmax(q·kᵀ / sqrt(head_dim) + causal mask) · v, with every head having its own k and v."""
+
+    def __init__(self, cfg: ModelConfig, layer_idx: int):
+        super().__init__(cfg, layer_idx)
+        if cfg.n_kv_heads != cfg.n_heads:
+            raise ValueError("mha needs n_kv_heads == n_heads; use 'gqa' for fewer KV heads")
+        d, h, hd = cfg.d_model, cfg.n_heads, cfg.head_dim
+        self.n_heads, self.head_dim = h, hd
+        # Separate projections with HF Llama names (export / TinyLlama loading).
+        self.q_proj = nn.Linear(d, h * hd, bias=False)
+        self.k_proj = nn.Linear(d, h * hd, bias=False)
+        self.v_proj = nn.Linear(d, h * hd, bias=False)
+        self.o_proj = nn.Linear(h * hd, d, bias=False)
+
+    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
+        """cache: a kv_cache.LayerKV. x then holds only the new tokens (positions cache.start..)."""
+        B, T, _ = x.shape
+        # (B, T, h*hd) -> (B, h, T, hd)
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        cos, sin = rope
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        mask = {"is_causal": True}
+        if cache is not None:
+            k, v = cache.update(k, v)                    # all cached positions, new ones included
+            mask = cache.attn_kwargs                     # the cache knows positions and padding
+        # Fused kernel (FlashAttention / memory-efficient), never materialises the T×T matrix.
+        y = F.scaled_dot_product_attention(q, k, v, **mask)
+        return self.o_proj(y.transpose(1, 2).reshape(B, T, -1))
+
+    def cache_bytes_per_token(self, dtype_bytes: int = 2) -> int:
+        return 2 * self.n_heads * self.head_dim * dtype_bytes  # one K and one V vector per head

src/tinyllm/config.py

@@ -0,0 +1,98 @@
+"""Model configuration and named presets (docs/ARCHITECTURE.md)."""
+
+from dataclasses import asdict, dataclass, field, replace
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    # Sizes
+    vocab_size: int = 8192
+    max_seq_len: int = 512
+    d_model: int = 512
+    n_layers: int = 8
+    n_heads: int = 8
+    n_kv_heads: int | None = None   # None = n_heads (plain MHA); fewer = GQA (Phase 4.9a)
+    ffn_hidden: int = 1408
+    # Numerics
+    norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    init_std: float = 0.02
+    tie_embeddings: bool = True
+    # Attention variant per layer, e.g. ("gdn", "gdn", "gdn", "mha") * 2 for a hybrid.
+    # None = all "mha", or all "gqa" when n_kv_heads < n_heads (4.9a).
+    layer_types: tuple[str, ...] | None = field(default=None)
+
+    def __post_init__(self) -> None:
+        # frozen dataclass: fill defaults through object.__setattr__
+        if self.n_kv_heads is None:
+            object.__setattr__(self, "n_kv_heads", self.n_heads)
+        if self.layer_types is None:
+            default = "gqa" if self.n_kv_heads < self.n_heads else "mha"
+            object.__setattr__(self, "layer_types", (default,) * self.n_layers)
+        else:
+            object.__setattr__(self, "layer_types", tuple(self.layer_types))
+
+        if self.d_model % self.n_heads:
+            raise ValueError(f"d_model {self.d_model} not divisible by n_heads {self.n_heads}")
+        if self.n_heads % self.n_kv_heads:
+            raise ValueError(f"n_heads {self.n_heads} not divisible by n_kv_heads {self.n_kv_heads}")
+        if self.head_dim % 2:
+            raise ValueError(f"head_dim {self.head_dim} must be even for RoPE")
+        if len(self.layer_types) != self.n_layers:
+            raise ValueError(f"{len(self.layer_types)} layer_types for {self.n_layers} layers")
+
+    @property
+    def head_dim(self) -> int:
+        return self.d_model // self.n_heads
+
+    def estimate_params(self) -> int:
+        """Parameter count for softmax-attention layers (mha / gqa).
+
+        embeddings V·d (tied with the LM head)
+        + per layer: q,o 2·d² + k,v 2·d·kv_dim + MLP 3·d·ffn + 2 RMSNorms 2·d
+        + final RMSNorm d
+        With n_kv_heads == n_heads this is the ARCHITECTURE.md formula
+        V·d + L·(4·d² + 3·d·ffn) + norms.
+        """
+        d, kv_dim = self.d_model, self.n_kv_heads * self.head_dim
+        per_layer = 2 * d * d + 2 * d * kv_dim + 3 * d * self.ffn_hidden + 2 * d
+        emb = self.vocab_size * d * (1 if self.tie_embeddings else 2)
+        return emb + self.n_layers * per_layer + d
+
+    def flops_per_token(self, seq_len: int) -> int:
+        """Training FLOPs per token (forward + backward), the PaLM MFU convention
+        (Chowdhery et al. 2022, App. B): 6·N + 12·L·(n_heads·head_dim)·T.
+
+        N counts only weights used in matmuls: per-layer projections plus the LM head (V·d).
+        The input embedding lookup and the RMSNorm weights do no matmul work. The attention
+        term counts the full T×T score matrix (no causal discount), as PaLM does, so numbers
+        compare with published MFU. Softmax-attention layers (mha / gqa) only.
+        """
+        d, kv_dim = self.d_model, self.n_kv_heads * self.head_dim
+        matmul = self.n_layers * (2 * d * d + 2 * d * kv_dim + 3 * d * self.ffn_hidden)
+        matmul += self.vocab_size * d
+        attn = 12 * self.n_layers * self.n_heads * self.head_dim * seq_len
+        return 6 * matmul + attn
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "ModelConfig":
+        return cls(**d)
+
+
+PRESETS: dict[str, ModelConfig] = {
+    "smoke": ModelConfig(d_model=384, n_layers=6, n_heads=6, ffn_hidden=1024),
+    "main": ModelConfig(d_model=512, n_layers=8, n_heads=8, ffn_hidden=1408),
+}
+
+
+def get_config(name: str, **overrides) -> ModelConfig:
+    """A preset with optional overrides, e.g. get_config("main", n_kv_heads=2)."""
+    cfg = PRESETS[name]
+    if overrides:
+        if ("n_layers" in overrides or "n_kv_heads" in overrides) and "layer_types" not in overrides:
+            overrides["layer_types"] = None  # re-derive for the new depth / head layout
+        cfg = replace(cfg, **overrides)
+    return cfg

src/tinyllm/cuda_graph.py

@@ -0,0 +1,304 @@
+"""CUDA-graph decode (4.11, docs/INFERENCE.md).
+
+A decode step of a small model is launch-bound: ~10 kernels per layer, each a few µs of
+GPU work, and the CPU needs about as long to launch each one. A CUDA graph records the
+whole step once and replays it with a single launch.
+
+A graph replays fixed kernels on fixed addresses, so the step must have static shapes:
+- StaticKVCache keeps the position as a device tensor. K/V are written with index_copy_
+  at that position, RoPE is gathered at it, and every decode step attends over a fixed
+  window of W keys, with keys past the position masked out.
+- W comes from length buckets (128, 256, ..., max_len). Each bucket has its own graph;
+  the step uses the smallest bucket that holds pos + 1 keys. All graphs share one memory
+  pool and the same static buffers (cache, input token, position).
+- Prefill stays eager: it runs once per request and has a different shape each time.
+  Sampling stays eager too.
+
+Not bit-identical to eager generate(): eager decode attends over exactly pos + 1 keys
+with no mask; the graph attends over W keys with a mask, and SDPA picks a different
+kernel for that. Graph replay *is* bit-identical to the same static step run eagerly.
+
+Autocast runs with cache_enabled=False: its weight-cast cache doesn't survive capture.
+"""
+
+import time
+
+import torch
+from torch import Tensor
+
+from tinyllm.kv_cache import causal_kwargs, padded_causal_mask
+from tinyllm.model import Transformer
+
+
+def default_buckets(max_len: int, first: int = 128) -> tuple[int, ...]:
+    """first, 2*first, ... below max_len, then max_len itself."""
+    out, b = [], first
+    while b < max_len:
+        out.append(b)
+        b *= 2
+    return (*out, max_len)
+
+
+class StaticKVCache:
+    """Shared-position cache with fixed shapes (B rows, max_len positions).
+
+    Same model protocol as kv_cache.KVCache. Two modes:
+    - prefill (eager, position 0): as KVCache, writes and returns positions 0..T-1.
+    - decode (`window` set, T=1): position from the device tensor `pos`; returns the first
+      `window` keys with a mask, so shapes don't depend on the position.
+    first_valid (B,) is the left-pad count per row; zeros when unpadded.
+    """
+
+    def __init__(self, n_layers: int, batch: int, max_len: int, device):
+        self.n_layers, self.batch, self.max_len = n_layers, batch, max_len
+        self.k: list[Tensor | None] = [None] * n_layers
+        self.v: list[Tensor | None] = [None] * n_layers
+        self.pos = torch.zeros(1, dtype=torch.long, device=device)      # device copy: graph input
+        self.first_valid = torch.zeros(batch, dtype=torch.long, device=device)
+        self.padded = False
+        self.n = 0                       # host copy of pos, for bounds and bucket choice (no sync)
+        self.window: int | None = None   # None = prefill mode
+        self.attn_kwargs: dict = {}
+
+    def reset(self, pad_lens: Tensor | None = None) -> None:
+        self.n, self.window = 0, None
+        self.pos.zero_()
+        self.padded = pad_lens is not None
+        self.first_valid.copy_(pad_lens if self.padded else torch.zeros_like(self.first_valid))
+
+    def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
+        if self.window is None:                          # eager prefill from position 0
+            if self.n != 0:
+                raise ValueError("StaticKVCache prefills only from position 0")
+            cos, sin = rope(T, 0)
+            self.attn_kwargs = ({"attn_mask": padded_causal_mask(0, T, self.first_valid)}
+                                if self.padded else causal_kwargs(0, T, cos.device))
+            return cos, sin
+        if T != 1:
+            raise ValueError("decode mode takes one token per row")
+        k_pos = torch.arange(self.window, device=self.pos.device)
+        valid = (k_pos <= self.pos) & (k_pos >= self.first_valid[:, None])        # (B, W)
+        self.attn_kwargs = {"attn_mask": valid[:, None, None, :]}                 # (B, 1, 1, W)
+        return rope.cos.index_select(0, self.pos), rope.sin.index_select(0, self.pos)
+
+    def end_step(self, T: int) -> None:
+        if self.window is None:
+            self.n = T
+            self.pos.fill_(T)
+        else:
+            self.pos.add_(1)             # captured into the graph; the host copy moves in step()
+
+    def layer(self, i: int) -> "StaticLayerKV":
+        return StaticLayerKV(self, i)
+
+    def nbytes(self) -> int:
+        return sum(t.numel() * t.element_size() for t in self.k + self.v if t is not None)
+
+
+class StaticLayerKV:
+    def __init__(self, cache: StaticKVCache, layer_idx: int):
+        self.cache = cache
+        self.i = layer_idx
+
+    @property
+    def attn_kwargs(self) -> dict:
+        return self.cache.attn_kwargs
+
+    def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        c, i = self.cache, self.i
+        B, H, T, D = k.shape
+        if c.k[i] is None:
+            shape = (c.batch, H, c.max_len, D)
+            c.k[i] = torch.zeros(shape, dtype=k.dtype, device=k.device)
+            c.v[i] = torch.zeros(shape, dtype=v.dtype, device=v.device)
+        elif c.k[i].dtype != k.dtype:
+            raise ValueError(f"cache holds {c.k[i].dtype}, got {k.dtype}")
+        if c.window is None:
+            if T > c.max_len:
+                raise ValueError(f"KV cache full: position {T} > max_len {c.max_len}")
+            c.k[i][:, :, :T] = k
+            c.v[i][:, :, :T] = v
+            return c.k[i][:, :, :T], c.v[i][:, :, :T]
+        c.k[i].index_copy_(2, c.pos, k)
+        c.v[i].index_copy_(2, c.pos, v)
+        return c.k[i][:, :, :c.window], c.v[i][:, :, :c.window]
+
+
+class GraphDecoder:
+    """Prefill eagerly, then decode one token per row with one CUDA-graph replay per step.
+
+    Reusable: prefill() resets the cache, and graphs captured for one request serve the
+    next (same batch size). use_graph=False runs the identical static step eagerly, the
+    reference that replay must match bit for bit.
+    """
+
+    def __init__(self, model: Transformer, batch: int, max_len: int | None = None,
+                 buckets: tuple[int, ...] | None = None, use_graph: bool = True):
+        device = next(model.parameters()).device
+        if device.type != "cuda":
+            raise ValueError("CUDA graphs need a CUDA model")
+        self.model = model
+        self.max_len = max_len or model.cfg.max_seq_len
+        if self.max_len > model.cfg.max_seq_len:
+            raise ValueError(f"max_len {self.max_len} > max_seq_len {model.cfg.max_seq_len}")
+        self.buckets = tuple(sorted(buckets)) if buckets else default_buckets(self.max_len)
+        if self.buckets[-1] < self.max_len:
+            raise ValueError(f"largest bucket {self.buckets[-1]} < max_len {self.max_len}")
+        self.cache = StaticKVCache(model.cfg.n_layers, batch, self.max_len, device)
+        self.tok = torch.zeros(batch, 1, dtype=torch.long, device=device)
+        self.use_graph = use_graph
+        self.graphs: dict[int, tuple[torch.cuda.CUDAGraph, Tensor]] = {}
+        self.pool = torch.cuda.graph_pool_handle()
+
+    def _autocast(self):
+        return torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=False)
+
+    @torch.no_grad()
+    def prefill(self, idx: Tensor, pad_lens: Tensor | None = None) -> Tensor:
+        """idx (B, T), left-padded by pad_lens (B,) if given. Returns logits (B, T, V)."""
+        if idx.shape[0] != self.cache.batch:
+            raise ValueError(f"batch {idx.shape[0]} != decoder batch {self.cache.batch}")
+        self.cache.reset(pad_lens)
+        with self._autocast():
+            logits, _ = self.model(idx, cache=self.cache)
+        return logits
+
+    def _forward(self) -> Tensor:
+        with self._autocast():
+            logits, _ = self.model(self.tok, cache=self.cache)
+        return logits
+
+    def _capture(self, window: int) -> tuple[torch.cuda.CUDAGraph, Tensor]:
+        # Warm up on a side stream (lazy cuBLAS / SDPA setup must not happen during capture).
+        # The warm-up runs the real step: its K/V writes are the ones replay will make, and
+        # the position it advances is put back.
+        saved = self.cache.pos.clone()
+        side = torch.cuda.Stream()
+        side.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(side):
+            self._forward()
+        torch.cuda.current_stream().wait_stream(side)
+        self.cache.pos.copy_(saved)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, pool=self.pool):
+            out = self._forward()
+        return graph, out
+
+    def bucket(self, n_keys: int) -> int:
+        return next(b for b in self.buckets if b >= n_keys)
+
+    @torch.no_grad()
+    def step(self, tok: Tensor) -> Tensor:
+        """tok (B, 1) → logits (B, 1, V). With graphs, the result is a static buffer that
+        the next step overwrites: read it (or sample from it) first."""
+        c = self.cache
+        if c.n >= self.max_len:
+            raise ValueError(f"KV cache full: position {c.n + 1} > max_len {self.max_len}")
+        c.window = self.bucket(c.n + 1)
+        self.tok.copy_(tok)
+        if not self.use_graph:
+            out = self._forward()
+        else:
+            if c.window not in self.graphs:
+                self.graphs[c.window] = self._capture(c.window)
+            graph, out = self.graphs[c.window]
+            graph.replay()
+        c.n += 1
+        return out
+
+
+class SlotGraphDecoder:
+    """The continuous engine's decode step (one token per SlotKVCache slot) as CUDA graphs
+    (4.11b). Admission, release and prefill stay eager in the engine.
+
+    Why it can be a graph: the batch is always n_slots rows (free slots compute too), and
+    slot lengths / pad starts / active flags live in fixed device tensors that admit() and
+    release() update in place. The one changing shape is the attention window, so the cache
+    gets a fixed window from the length buckets, with keys past each slot's length masked.
+
+    Each step writes the token into a static buffer and replays the graph for the smallest
+    bucket that holds the longest slot + 1. The replay advances the device lengths; the host
+    lengths move in advance_host(). use_graph=False runs the identical windowed step eagerly,
+    the reference that replay must match bit for bit.
+    """
+
+    def __init__(self, model: Transformer, slots, buckets: tuple[int, ...] | None = None,
+                 use_graph: bool = True):
+        device = next(model.parameters()).device
+        if device.type != "cuda":
+            raise ValueError("CUDA graphs need a CUDA model")
+        if not hasattr(slots, "advance_host"):
+            raise ValueError("SlotGraphDecoder supports the slot cache (kv='slots') only")
+        self.model, self.slots = model, slots
+        self.buckets = tuple(sorted(buckets)) if buckets else default_buckets(slots.max_len)
+        if self.buckets[-1] < slots.max_len:
+            raise ValueError(f"largest bucket {self.buckets[-1]} < max_len {slots.max_len}")
+        self.tok = torch.zeros(slots.n_slots, 1, dtype=torch.long, device=device)
+        self.use_graph = use_graph
+        self.graphs: dict[int, tuple[torch.cuda.CUDAGraph, Tensor]] = {}
+        self.pool = torch.cuda.graph_pool_handle()
+
+    def _forward(self) -> Tensor:
+        with torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=False):
+            logits, _ = self.model(self.tok, cache=self.slots)
+        return logits
+
+    def _capture(self) -> tuple[torch.cuda.CUDAGraph, Tensor]:
+        # The warm-up runs the real step (its K/V writes are the ones replay will make) and
+        # advances the lengths; capture runs the host half again. Put both back each time.
+        c = self.slots
+        lens, lens_dev = list(c.lens), c._lens_dev.clone()
+        side = torch.cuda.Stream()
+        side.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(side):
+            self._forward()
+        torch.cuda.current_stream().wait_stream(side)
+        c.lens[:] = lens
+        c._lens_dev.copy_(lens_dev)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, pool=self.pool):
+            out = self._forward()
+        c.lens[:] = lens
+        return graph, out
+
+    def bucket(self, n_keys: int) -> int:
+        return next(b for b in self.buckets if b >= n_keys)
+
+    @torch.no_grad()
+    def capture(self, windows: tuple[int, ...] | None = None) -> float:
+        """Capture graphs ahead of time (default: every bucket), so no timed step pays for it.
+        Safe between steps: the warm-up's K/V writes land at each slot's next position, which
+        the next real step overwrites. Returns the seconds it took."""
+        t0 = time.perf_counter()
+        c = self.slots
+        try:
+            for w in windows or self.buckets:
+                if w not in self.graphs:
+                    c.window = w
+                    self.graphs[w] = self._capture()
+        finally:
+            c.window = None
+        torch.cuda.synchronize()
+        return time.perf_counter() - t0
+
+    @torch.no_grad()
+    def step(self, tok: Tensor) -> Tensor:
+        """tok (n_slots, 1) → logits (n_slots, 1, V). With graphs, the result is a static
+        buffer that the next step overwrites: sample from it first."""
+        c = self.slots
+        need = max(c.lens) + 1
+        if need > c.max_len:
+            raise ValueError(f"a slot is full (max_len {c.max_len})")
+        c.window = self.bucket(need)
+        self.tok.copy_(tok)
+        try:
+            if not self.use_graph:
+                return self._forward()
+            if c.window not in self.graphs:
+                self.graphs[c.window] = self._capture()
+            graph, out = self.graphs[c.window]
+            graph.replay()
+            c.advance_host()
+            return out
+        finally:
+            c.window = None

src/tinyllm/data.py

@@ -0,0 +1,130 @@
+"""Dataset loading and text cleanup (D-008, D-009, D-018).
+
+Every step that reads TinyStories (tokenizer training, tokenization) goes
+through `load_clean`, so they all see exactly the same text.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from datasets import Dataset, load_dataset
+
+DATASET = "roneneldan/TinyStories"
+DATA_DIR = Path(__file__).resolve().parents[2] / "data"  # git-ignored
+TOKEN_DTYPE = np.uint16  # vocab 8192 < 65536 (D-009)
+
+# D-018: TinyStories has UTF-8 punctuation that was decoded as Windows-1252
+# ("mojibake"), e.g. “ -> "“". Map it back to plain ASCII.
+# Order matters: the 3-character sequences must run before the bare "â€".
+_REPLACEMENTS: list[tuple[str, str]] = [
+    ("“", '"'),    # “ left double quote
+    ("â€\x9d", '"'), # ” right double quote (when the 3rd byte survived)
+    ("’", "'"),    # ’ apostrophe / right single quote
+    ("‘", "'"),    # ‘ left single quote
+    ("–", "-"),    # – en dash
+    ("—", "-"),    # — em dash
+    ("…", "..."),  # … ellipsis
+    ("â€", '"'),     # ” right double quote whose 3rd byte (0x9D) was lost
+    ("Â\xa0", " "),  # no-break space
+    ("Â\xad", ""),   # soft hyphen
+    # Real (correctly encoded) typographic punctuation -> ASCII.
+    ("“", '"'), ("”", '"'), ("‘", "'"), ("’", "'"),
+    ("–", "-"), ("—", "-"), ("…", "..."), ("\xa0", " "), ("\xad", ""),
+]
+
+
+def clean_text(text: str) -> str | None:
+    """Return cleaned ASCII text, or None if the story should be dropped.
+
+    Dropped: empty stories, and stories that still contain non-ASCII
+    characters after the replacements (rare leftovers such as 'é').
+    """
+    if not text.isascii():
+        for bad, good in _REPLACEMENTS:
+            if bad in text:
+                text = text.replace(bad, good)
+    text = text.strip()
+    if not text or not text.isascii():
+        return None
+    return text
+
+
+def _clean_batch(batch: dict) -> dict:
+    return {"text": [clean_text(t) or "" for t in batch["text"]]}
+
+
+def _keep_batch(batch: dict) -> list[bool]:
+    return [t != "" for t in batch["text"]]
+
+
+def load_clean(split: str, num_proc: int = 4) -> Dataset:
+    """Load a TinyStories split ('train' or 'validation') with D-018 cleanup applied.
+
+    `datasets` caches the result, so only the first call does the work.
+    """
+    ds = load_dataset(DATASET, split=split)
+    # With split=... (and no streaming) this is always a Dataset. The check narrows
+    # the declared Union return type for type checkers, and fails loudly otherwise.
+    assert isinstance(ds, Dataset), f"expected Dataset, got {type(ds).__name__}"
+    ds = ds.map(_clean_batch, batched=True, num_proc=num_proc, desc=f"clean {split}")
+    return ds.filter(_keep_batch, batched=True, num_proc=num_proc, desc=f"drop {split}")
+
+
+def bin_path(split: str) -> Path:
+    """'train' -> data/train.bin, 'validation' -> data/val.bin"""
+    return DATA_DIR / f"{'val' if split == 'validation' else split}.bin"
+
+
+def write_tokens(split: str, tok, eos_id: int, batch_size: int = 10_000) -> tuple[int, int]:
+    """Encode a cleaned split into one flat uint16 file: story EOS story EOS ...
+
+    Streams batch by batch, so RAM use stays small. Writes to a .tmp file and
+    renames at the end, so a crash never leaves a half-written .bin behind.
+    Returns (n_stories, n_tokens).
+    """
+    ds = load_clean(split)
+    out = bin_path(split)
+    out.parent.mkdir(exist_ok=True)
+    tmp = out.with_suffix(".tmp")
+    n_tokens = 0
+    with open(tmp, "wb") as f:
+        for i in range(0, len(ds), batch_size):
+            encs = tok.encode_batch(ds[i: i + batch_size]["text"])  # Rust, multi-threaded
+            ids = np.fromiter(
+                (t for e in encs for t in (*e.ids, eos_id)), dtype=np.int64
+            )
+            assert ids.max() < np.iinfo(TOKEN_DTYPE).max
+            ids.astype(TOKEN_DTYPE).tofile(f)
+            n_tokens += len(ids)
+    tmp.replace(out)
+    return len(ds), n_tokens
+
+
+def load_tokens(split: str) -> np.memmap:
+    """Memory-map a .bin file: no RAM cost, the OS pages it in on demand."""
+    return np.memmap(bin_path(split), dtype=TOKEN_DTYPE, mode="r")
+
+
+def get_batch(
+    tokens: np.ndarray,
+    batch_size: int,
+    seq_len: int,
+    rng: np.random.Generator,
+    device: str | torch.device = "cpu",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Random training windows: x = tokens[s : s+T], y = tokens[s+1 : s+T+1].
+
+    Windows start anywhere and cross story boundaries; EOS tells the model where a
+    story ends. Sampling with replacement from a seeded `rng` makes runs reproducible,
+    and saving `rng.bit_generator.state` in a checkpoint makes resume exact.
+    On CUDA the batch goes through pinned memory so the copy can overlap with compute.
+    """
+    if len(tokens) < seq_len + 1:
+        raise ValueError(f"need at least {seq_len + 1} tokens, got {len(tokens)}")
+    starts = rng.integers(0, len(tokens) - seq_len, size=batch_size)    # last valid start: len-T-1
+    idx = starts[:, None] + np.arange(seq_len + 1)                       # (B, T+1)
+    chunk = torch.from_numpy(tokens[idx].astype(np.int64))               # gather, then widen for embedding
+    if torch.device(device).type == "cuda":
+        chunk = chunk.pin_memory().to(device, non_blocking=True)
+    return chunk[:, :-1], chunk[:, 1:]

src/tinyllm/engine.py

@@ -0,0 +1,431 @@
+"""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)

src/tinyllm/generate.py

@@ -0,0 +1,140 @@
+"""Text generation.
+
+- sample(): naive, a full forward pass over the context per new token (no cache).
+  Used for training samples, and the reference that every faster path must match (4.1).
+- generate(): KV-cached, one prefill pass then one new token per forward (4.3).
+  cuda_graph=True replays each decode step as one CUDA graph (4.11, tinyllm.cuda_graph).
+Both choose tokens with sample_next (4.2).
+"""
+
+import torch
+from torch import Tensor
+
+from tinyllm.kv_cache import KVCache
+from tinyllm.model import Transformer
+
+
+def sample_next(
+    logits: Tensor,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    top_p: float | None = None,
+    generator: torch.Generator | None = None,
+) -> Tensor:
+    """Pick the next token from last-position logits (B, V) → (B, 1). Step 4.2.
+
+    Order (same as HF's warpers): temperature, then top-k, then top-p.
+    - temperature 0 means greedy (argmax); top_k / top_p are then irrelevant.
+    - top_k keeps the k most likely tokens.
+    - top_p (nucleus) keeps the smallest set of most likely tokens whose total
+      probability is >= top_p. The most likely token is always kept.
+    """
+    if top_k is not None and top_k < 1:
+        raise ValueError(f"top_k must be >= 1, got {top_k}")
+    if top_p is not None and not 0.0 < top_p <= 1.0:
+        raise ValueError(f"top_p must be in (0, 1], got {top_p}")
+    logits = logits.float()
+    if temperature == 0:
+        return logits.argmax(-1, keepdim=True)
+    logits = logits / temperature
+    if top_k is not None:
+        kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+    if top_p is not None and top_p < 1.0:
+        sorted_logits, order = logits.sort(-1, descending=True)
+        probs = sorted_logits.softmax(-1)
+        mass_before = probs.cumsum(-1) - probs           # probability of all strictly better tokens
+        drop_sorted = mass_before >= top_p               # first token has mass_before 0: always kept
+        drop = drop_sorted.scatter(-1, order, drop_sorted)
+        logits = logits.masked_fill(drop, float("-inf"))
+    return torch.multinomial(logits.softmax(-1), 1, generator=generator)
+
+
+@torch.no_grad()
+def sample(
+    model: Transformer,
+    idx: Tensor,
+    max_new_tokens: int,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    eos_id: int | None = None,
+    generator: torch.Generator | None = None,
+    top_p: float | None = None,
+) -> Tensor:
+    """Append up to max_new_tokens to idx (1, T). Stops early when eos_id is sampled.
+
+    Token choice is sample_next (temperature, top_k, top_p).
+    """
+    was_training = model.training
+    model.eval()
+    max_len = model.cfg.max_seq_len
+    for _ in range(max_new_tokens):
+        context = idx[:, -max_len:]  # no cache: the window slides once the context is full
+        with torch.autocast(idx.device.type, dtype=torch.bfloat16, enabled=idx.device.type == "cuda"):
+            logits, _ = model(context)
+        nxt = sample_next(logits[:, -1, :], temperature, top_k, top_p, generator)
+        idx = torch.cat([idx, nxt], dim=1)
+        if eos_id is not None and nxt.item() == eos_id:
+            break
+    model.train(was_training)
+    return idx
+
+
+@torch.no_grad()
+def generate(
+    model: Transformer,
+    idx: Tensor,
+    max_new_tokens: int,
+    temperature: float = 1.0,
+    top_k: int | None = None,
+    top_p: float | None = None,
+    eos_id: int | None = None,
+    generator: torch.Generator | None = None,
+    cuda_graph=False,
+) -> Tensor:
+    """KV-cached generation. Append up to max_new_tokens to idx (B, T).
+
+    Prefill runs the prompt once and fills the cache; each decode step feeds only the
+    newest token. Stops early once every row has produced eos_id (rows that finish
+    first keep generating; the caller trims after EOS).
+
+    Unlike sample(), the context can't slide: RoPE positions are baked into the cached
+    keys. So prompt + max_new_tokens must fit in max_seq_len.
+
+    cuda_graph: False (eager), True (build a GraphDecoder), or a GraphDecoder to reuse,
+    so the graphs it captured serve several calls with the same batch size (4.11).
+    """
+    B, T = idx.shape
+    max_len = model.cfg.max_seq_len
+    if T + max_new_tokens > max_len:
+        raise ValueError(f"prompt {T} + new {max_new_tokens} exceeds max_seq_len {max_len}")
+    was_training = model.training
+    model.eval()
+    if cuda_graph:
+        from tinyllm.cuda_graph import GraphDecoder
+        dec = cuda_graph if isinstance(cuda_graph, GraphDecoder) else GraphDecoder(model, B)
+        prefill, decode = dec.prefill, dec.step
+    else:
+        cache = KVCache(model.cfg.n_layers, max_len)
+        autocast = dict(device_type=idx.device.type, dtype=torch.bfloat16,
+                        enabled=idx.device.type == "cuda")
+
+        def prefill(x: Tensor) -> Tensor:
+            with torch.autocast(**autocast):
+                return model(x, cache=cache)[0]
+        decode = prefill
+    done = torch.zeros(B, dtype=torch.bool, device=idx.device)
+    out = [idx]
+    logits = prefill(idx)
+    for i in range(max_new_tokens):
+        nxt = sample_next(logits[:, -1, :], temperature, top_k, top_p, generator)
+        out.append(nxt)
+        if eos_id is not None:
+            done |= nxt.squeeze(-1) == eos_id
+            if done.all():
+                break
+        if i + 1 == max_new_tokens:
+            break                                              # no forward for a token we won't use
+        logits = decode(nxt)                                   # one token per row
+    model.train(was_training)
+    return torch.cat(out, dim=1)

src/tinyllm/hf_llama.py

@@ -0,0 +1,192 @@
+"""Load a pretrained Hugging Face Llama checkpoint into our Transformer (step 4.9b, D-015 / D-026).
+
+Only our code runs the model. Hugging Face pieces used here are file formats and downloads:
+`huggingface_hub` fetches the files, our `safetensors_io` reads the weights, and `tokenizers`
+reads `tokenizer.json`. HF Transformers is used only as the logits reference in tests (4.9c).
+
+Our parameter names already follow HF Llama (embed_tokens, layers.N.self_attn.q_proj, ...,
+input_layernorm, mlp.gate_proj, norm, lm_head), so the mapping just drops HF's "model." prefix.
+RoPE uses the same rotate_half layout as HF, so q/k weights need no permutation.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from tinyllm.config import ModelConfig
+from tinyllm.model import RotaryEmbedding, Transformer
+from tinyllm.safetensors_io import load_file
+
+TINYLLAMA = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+TINYLLAMA_REV = "fe8a4ea1ffedaf415f4da2f062534de366a451e6"   # pinned (D-026)
+FILES = ["config.json", "generation_config.json", "model.safetensors",
+         "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]
+
+
+def config_from_hf(hf: dict) -> ModelConfig:
+    """HF LlamaConfig dict -> ModelConfig. Raises on any feature our model doesn't implement."""
+    unsupported = []
+    if hf.get("model_type") != "llama":
+        unsupported.append(f"model_type={hf.get('model_type')!r}")
+    if hf.get("hidden_act", "silu") != "silu":
+        unsupported.append(f"hidden_act={hf['hidden_act']!r}")
+    if hf.get("rope_scaling"):
+        unsupported.append(f"rope_scaling={hf['rope_scaling']!r}")
+    # transformers >= 5 writes {"rope_parameters": {"rope_theta": ..., "rope_type": ...}} instead of
+    # a top-level rope_theta; scaled variants (linear, llama3, yarn, ...) are a rope_type there.
+    rp = hf.get("rope_parameters") or {}
+    if rp.get("rope_type", "default") != "default":
+        unsupported.append(f"rope_parameters={rp!r}")
+    rope_theta = rp.get("rope_theta", hf.get("rope_theta", 10000.0))
+    if hf.get("attention_bias") or hf.get("mlp_bias"):
+        unsupported.append("bias in attention / MLP")
+    if hf.get("pretraining_tp", 1) != 1:
+        unsupported.append(f"pretraining_tp={hf['pretraining_tp']}")
+    hd = hf.get("head_dim")
+    if hd is not None and hd != hf["hidden_size"] // hf["num_attention_heads"]:
+        unsupported.append(f"head_dim={hd} != hidden_size / num_attention_heads")
+    if unsupported:
+        raise ValueError("unsupported Llama config: " + ", ".join(unsupported))
+    return ModelConfig(
+        vocab_size=hf["vocab_size"],
+        max_seq_len=hf["max_position_embeddings"],
+        d_model=hf["hidden_size"],
+        n_layers=hf["num_hidden_layers"],
+        n_heads=hf["num_attention_heads"],
+        n_kv_heads=hf.get("num_key_value_heads", hf["num_attention_heads"]),
+        ffn_hidden=hf["intermediate_size"],
+        norm_eps=hf["rms_norm_eps"],
+        rope_theta=rope_theta,
+        init_std=hf.get("initializer_range", 0.02),
+        tie_embeddings=hf.get("tie_word_embeddings", False),
+    )
+
+
+def hf_to_ours(name: str) -> str:
+    return name.removeprefix("model.")
+
+
+def ours_to_hf(name: str) -> str:
+    return name if name == "lm_head.weight" else "model." + name
+
+
+def download(repo: str = TINYLLAMA, revision: str = TINYLLAMA_REV, local_only: bool = False) -> Path:
+    """Snapshot directory in the HF cache (~/.cache/huggingface). Weights never go into the repo."""
+    from huggingface_hub import snapshot_download
+    return Path(snapshot_download(repo, revision=revision, allow_patterns=FILES, local_files_only=local_only))
+
+
+def load_from_dir(path: str | Path, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> Transformer:
+    """Build our Transformer from <path>/config.json and <path>/model.safetensors, strict.
+
+    The model is created on the meta device, so no fp32 copy is ever allocated (TinyLlama: 4.4 GB);
+    each weight is copied once from the memory-mapped file to `device` in `dtype` (D-023: bf16).
+    """
+    path = Path(path)
+    cfg = config_from_hf(json.loads((path / "config.json").read_text()))
+    state = {hf_to_ours(k): v for k, v in load_file(path / "model.safetensors", device, dtype).items()}
+    # Rotary inv_freq is a buffer in some HF exports; ours are recomputed and never stored.
+    state = {k: v for k, v in state.items() if not k.endswith("rotary_emb.inv_freq")}
+    if cfg.tie_embeddings:
+        state.pop("lm_head.weight", None)        # tied: the embedding is the head
+    with torch.device("meta"):
+        model = Transformer(cfg)
+    expected = {k: v for k, v in model.state_dict().items() if not (cfg.tie_embeddings and k == "lm_head.weight")}
+    missing, unexpected = sorted(set(expected) - set(state)), sorted(set(state) - set(expected))
+    bad_shape = [k for k in expected.keys() & state.keys() if tuple(expected[k].shape) != tuple(state[k].shape)]
+    if missing or unexpected or bad_shape:
+        raise ValueError(f"checkpoint mismatch: missing {missing[:5]}, unexpected {unexpected[:5]}, "
+                         f"shape {bad_shape[:5]}")
+    model.load_state_dict(state, strict=False, assign=True)   # strictness checked above (tie aware)
+    if cfg.tie_embeddings:
+        model.lm_head.weight = model.embed_tokens.weight
+    # Non-persistent buffers were created on meta: rebuild the RoPE tables on the real device.
+    model.rope = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta).to(device)
+    return model.eval()
+
+
+def load_hf_llama(repo: str = TINYLLAMA, revision: str = TINYLLAMA_REV, device: str = "cuda",
+                  dtype: torch.dtype = torch.bfloat16, local_only: bool = False) -> Transformer:
+    return load_from_dir(download(repo, revision, local_only), device, dtype)
+
+
+def save_as_hf(model: Transformer, path: str | Path, eos_id: int | None = None, bos_id: int | None = None) -> None:
+    """Write our model as an HF-style Llama directory (config.json, generation_config.json,
+    model.safetensors). Used by tests and 4.8b.
+
+    The token ids are written explicitly, even as null: a config without them gets LlamaConfig's
+    defaults (bos 1, eos 2), and HF generate() would then stop on the wrong token.
+    """
+    from tinyllm.safetensors_io import save_file
+    c = model.cfg
+    other = sorted(set(c.layer_types) - {"mha", "gqa"})
+    if other:
+        raise ValueError(f"only mha / gqa layers map to Llama; this model has {other}")
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=True)
+    sd = model.state_dict()
+    if c.tie_embeddings:
+        sd.pop("lm_head.weight")
+    dtype = str(next(iter(sd.values())).dtype).removeprefix("torch.")
+    ids = {"bos_token_id": bos_id, "eos_token_id": eos_id, "pad_token_id": None}
+    hf = {"architectures": ["LlamaForCausalLM"], "model_type": "llama", "hidden_act": "silu",
+          "vocab_size": c.vocab_size, "max_position_embeddings": c.max_seq_len, "hidden_size": c.d_model,
+          "num_hidden_layers": c.n_layers, "num_attention_heads": c.n_heads,
+          "num_key_value_heads": c.n_kv_heads, "intermediate_size": c.ffn_hidden, "rms_norm_eps": c.norm_eps,
+          "rope_theta": c.rope_theta, "initializer_range": c.init_std, "tie_word_embeddings": c.tie_embeddings,
+          "attention_bias": False, "rope_scaling": None, "dtype": dtype, **ids}
+    (path / "config.json").write_text(json.dumps(hf, indent=2))
+    # Padding a batch with EOS (GPT-2 style); our model has no pad token.
+    (path / "generation_config.json").write_text(json.dumps({**ids, "pad_token_id": eos_id}, indent=2))
+    save_file({ours_to_hf(k): v for k, v in sd.items()}, path / "model.safetensors", {"format": "pt"})
+
+
+def save_tokenizer_as_hf(path: str | Path, model_max_length: int | None = None) -> None:
+    """Our TinyStories BPE as an HF tokenizer directory (4.8b).
+
+    It was trained with the `tokenizers` library (D-007), so the file already is a tokenizer.json;
+    tokenizer_config.json only tells `AutoTokenizer` which class to use and which token is EOS.
+    Like our server: no BOS, nothing added on encode (the file has no post-processor), and no
+    "clean up" of spaces on decode, so decode(encode(text)) == text.
+    """
+    import shutil
+
+    from tinyllm.tokenizer import EOS, TOKENIZER_PATH
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(TOKENIZER_PATH, path / "tokenizer.json")
+    cfg = {"tokenizer_class": "PreTrainedTokenizerFast", "eos_token": EOS, "bos_token": None,
+           "unk_token": None, "pad_token": EOS, "add_bos_token": False, "add_eos_token": False,
+           "clean_up_tokenization_spaces": False, "model_max_length": model_max_length}
+    (path / "tokenizer_config.json").write_text(json.dumps(cfg, indent=2))
+
+
+# --- tokenizer -------------------------------------------------------------------------------------
+
+class HFTokenizer:
+    """tokenizer.json via the `tokenizers` library. encode() adds BOS (<s>) like the Llama tokenizer."""
+
+    def __init__(self, path: str | Path):
+        from tokenizers import Tokenizer
+        path = Path(path)
+        self.tok = Tokenizer.from_file(str(path / "tokenizer.json"))
+        cfg = json.loads((path / "tokenizer_config.json").read_text()) if (path / "tokenizer_config.json").exists() else {}
+        self.chat_template = cfg.get("chat_template")
+        self.bos_id = self.tok.token_to_id("<s>")
+        self.eos_id = self.tok.token_to_id("</s>")
+
+    def encode(self, text: str, bos: bool = True) -> list[int]:
+        ids = self.tok.encode(text, add_special_tokens=False).ids
+        return ([self.bos_id] if bos else []) + ids
+
+    def decode(self, ids: list[int]) -> str:
+        return self.tok.decode(ids, skip_special_tokens=True)
+
+    @staticmethod
+    def chat_prompt(user: str, system: str | None = None) -> str:
+        """TinyLlama-Chat v1.0 (Zephyr) template, with the assistant turn opened for generation."""
+        sys = f"<|system|>\n{system}</s>\n" if system else ""
+        return f"{sys}<|user|>\n{user}</s>\n<|assistant|>\n"

src/tinyllm/kernels/__init__.py

@@ -0,0 +1,4 @@
+"""Triton kernels (PLAN 4.10, D-016): each has a PyTorch reference it must match (`allclose`).
+
+Importing a kernel module needs Triton (Linux); callers fall back to PyTorch without it.
+"""

src/tinyllm/kernels/int8_gemv.py

@@ -0,0 +1,75 @@
+"""Step 4.10b (D-016): INT8 weight-only matmul for decode, reading the int8 weights directly.
+
+The PyTorch path in tinyllm.quant (Int8Linear, backend "torch") runs three kernels per layer:
+cast the int8 weight to bf16 (reads 1 byte and writes 2 per weight), matmul (reads those 2 bytes
+again), then scale. Here one kernel reads each int8 weight once, converts it in registers (exact:
+|q| <= 127), accumulates in fp32 and applies the per-row scale at the end:
+
+    y[m, n] = scale[n] * sum_k x[m, k] * q[n, k]          x: (M, K) bf16/fp16, q: (N, K) int8
+
+At decode, M (tokens in the step = batch) is small, so the step is bound by reading q. The tile
+over M is padded to 16 rows for tl.dot (tensor cores); the wasted flops are cheap next to the
+bytes. Each program owns BLOCK_N output columns and loops over K. Block sizes are autotuned per
+(N, K, M bucket) on first use, which must happen outside CUDA-graph capture (GraphDecoder warms up
+first). M above MAX_M (prefill) stays on the PyTorch path; see Int8Linear.forward.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+MAX_M = 64  # rows per call this kernel takes; more (prefill) is a GEMM, left to cuBLAS
+
+_CONFIGS = [triton.Config({"BLOCK_N": bn, "BLOCK_K": bk}, num_warps=w, num_stages=s)
+            for bn in (16, 32, 64) for bk in (128, 256) for w in (4,) for s in (3, 4)]
+_CONFIGS += [triton.Config({"BLOCK_N": 32, "BLOCK_K": 512}, num_warps=8, num_stages=3),
+             triton.Config({"BLOCK_N": 64, "BLOCK_K": 128}, num_warps=2, num_stages=4)]
+
+
+@triton.autotune(configs=_CONFIGS, key=["N", "K", "BLOCK_M"])
+@triton.jit
+def _int8_matmul_kernel(x_ptr, q_ptr, s_ptr, y_ptr, M, N, K, stride_xm, stride_qn, stride_ym,
+                        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    m_ok, n_ok = offs_m < M, offs_n < N
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        offs_k = k0 + tl.arange(0, BLOCK_K)
+        k_ok = offs_k < K
+        x = tl.load(x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+                    mask=m_ok[:, None] & k_ok[None, :], other=0.0)
+        q = tl.load(q_ptr + offs_n[:, None] * stride_qn + offs_k[None, :],
+                    mask=n_ok[:, None] & k_ok[None, :], other=0)
+        acc = tl.dot(x, tl.trans(q.to(x.dtype)), acc)       # int8 -> bf16 is exact
+    s = tl.load(s_ptr + offs_n, mask=n_ok, other=0.0)
+    y = acc * s[None, :]
+    tl.store(y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :], y, mask=m_ok[:, None] & n_ok[None, :])
+
+
+def int8_matmul(x: Tensor, q: Tensor, scale: Tensor) -> Tensor:
+    """x (M, K) bf16/fp16 CUDA, M <= MAX_M; q (N, K) int8; scale (N,) fp32 -> y (M, N) fp32.
+
+    Same values as F.linear(x, q.to(x.dtype)) * scale, except that the matmul result isn't rounded
+    to bf16 before scaling (closer to the fp32 dequantize-then-matmul reference, not further)."""
+    M, K = x.shape
+    N = q.shape[0]
+    if M > MAX_M or q.shape[1] != K or x.dtype not in (torch.bfloat16, torch.float16):
+        raise ValueError(f"int8_matmul: x {tuple(x.shape)} {x.dtype}, q {tuple(q.shape)}; needs M <= {MAX_M}, "
+                         "matching K, bf16/fp16 x")
+    x = x if x.stride(1) == 1 else x.contiguous()
+    q = q if q.stride(1) == 1 else q.contiguous()
+    y = torch.empty((M, N), device=x.device, dtype=torch.float32)
+    block_m = max(16, triton.next_power_of_2(M))
+    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]),)  # noqa: E731
+    _int8_matmul_kernel[grid](x, q, scale, y, M, N, K, x.stride(0), q.stride(0), y.stride(0), BLOCK_M=block_m)
+    return y
+
+
+def best_configs() -> dict[str, str]:
+    """The autotuned config per (N, K, BLOCK_M) key so far, for the benchmark JSON."""
+    return {str(k): str(v) for k, v in _int8_matmul_kernel.cache.items()}

src/tinyllm/kernels/rmsnorm.py

@@ -0,0 +1,88 @@
+"""Step 4.10a (D-016): RMSNorm, and residual add + RMSNorm, each as one Triton kernel.
+
+The PyTorch RMSNorm (tinyllm.model.RMSNorm, backend "torch") runs about seven small kernels per
+call: cast to fp32, square, mean, add eps + rsqrt, multiply, cast back, multiply by the weight. At
+decode a row is one token's 2,048 activations (4 KB in bf16), so each of those kernels is mostly
+fixed cost. Here one program per row loads the row once and does it all in registers:
+
+    rmsnorm:      y = (x * rsqrt(mean(x^2) + eps)).to(x.dtype) * w
+    add_rmsnorm:  h = x + a  (written out: the new residual stream);  y = rmsnorm(h)
+
+The roundings match the reference: h is rounded to its dtype before the statistic, the normalized
+row is rounded to x's dtype before the weight multiply, and the product is rounded to the promoted
+dtype of (x, w). Only the order of the fp32 sum differs, so results are allclose, not bitwise.
+
+One compiled config per (row width, dtypes): no autotuning, so nothing has to run before CUDA-graph
+capture beyond the first call (GraphDecoder warms up first).
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+
+MAX_D = 65536  # a whole row lives in one program's registers
+
+
+@triton.jit
+def _rmsnorm_kernel(x_ptr, a_ptr, h_ptr, w_ptr, y_ptr, D, eps, stride_x, stride_a, stride_h, stride_y,
+                    ADD: tl.constexpr, BLOCK_D: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK_D)
+    ok = offs < D
+    x = tl.load(x_ptr + row * stride_x + offs, mask=ok, other=0.0)
+    if ADD:
+        a = tl.load(a_ptr + row * stride_a + offs, mask=ok, other=0.0)
+        h = (x.to(tl.float32) + a.to(tl.float32)).to(h_ptr.dtype.element_ty)   # x + a in the promoted dtype
+        tl.store(h_ptr + row * stride_h + offs, h, mask=ok)
+        x = h
+    xf = x.to(tl.float32)
+    rstd = 1.0 / tl.sqrt(tl.sum(xf * xf, axis=0) / D + eps)
+    n = (xf * rstd).to(x.dtype).to(tl.float32)                                 # .to(x.dtype), as the reference
+    w = tl.load(w_ptr + offs, mask=ok, other=0.0).to(tl.float32)
+    tl.store(y_ptr + row * stride_y + offs, (n * w).to(y_ptr.dtype.element_ty), mask=ok)
+
+
+def _rows(t: Tensor) -> Tensor:
+    t = t.reshape(-1, t.shape[-1])
+    return t if t.stride(-1) == 1 else t.contiguous()
+
+
+def _launch(x: Tensor, a: Tensor | None, weight: Tensor, eps: float) -> tuple[Tensor, Tensor | None]:
+    D = x.shape[-1]
+    if not (x.is_cuda and weight.is_cuda) or weight.shape != (D,) or D > MAX_D:
+        raise ValueError(f"rmsnorm: x {tuple(x.shape)} on {x.device}, weight {tuple(weight.shape)}; "
+                         f"needs CUDA tensors, weight of shape ({D},), width <= {MAX_D}")
+    x2 = _rows(x)
+    h = h2 = a2 = None
+    if a is not None:
+        if a.shape != x.shape:
+            raise ValueError(f"add_rmsnorm: x {tuple(x.shape)} and a {tuple(a.shape)} differ")
+        a2 = _rows(a)
+        h = torch.empty(x.shape, device=x.device, dtype=torch.result_type(x, a))
+        h2 = h.view(-1, D)
+    src = x2 if h2 is None else h2
+    y = torch.empty(x.shape, device=x.device, dtype=torch.result_type(src, weight))
+    y2 = y.view(-1, D)
+    rows = x2.shape[0]
+    if rows:
+        block = triton.next_power_of_2(D)
+        warps = 4 if block <= 2048 else 8 if block <= 8192 else 16
+        _rmsnorm_kernel[(rows,)](
+            x2, x2 if a2 is None else a2, x2 if h2 is None else h2, weight, y2, D, eps,
+            x2.stride(0), 0 if a2 is None else a2.stride(0), 0 if h2 is None else h2.stride(0), y2.stride(0),
+            ADD=a is not None, BLOCK_D=block, num_warps=warps)
+    return y, h
+
+
+def rmsnorm(x: Tensor, weight: Tensor, eps: float) -> Tensor:
+    """x (..., D) CUDA -> y (..., D) in result_type(x, weight); same values as tinyllm.model.RMSNorm."""
+    return _launch(x, None, weight, eps)[0]
+
+
+def add_rmsnorm(x: Tensor, a: Tensor, weight: Tensor, eps: float) -> tuple[Tensor, Tensor]:
+    """(y, h) with h = x + a and y = rmsnorm(h, weight, eps), from one pass over the row."""
+    y, h = _launch(x, a, weight, eps)
+    return y, h

src/tinyllm/kv_cache.py

@@ -0,0 +1,238 @@
+"""KV caches for incremental decoding (docs/INFERENCE.md).
+
+The model talks to a cache through three calls per forward pass:
+    rope = cache.begin_step(T, rope_module)   # positions of the T new tokens; builds the mask
+    view = cache.layer(i)                     # per layer: view.update(k, v), view.attn_kwargs
+    cache.end_step(T)                         # the new tokens are now stored
+
+- KVCache (4.3, 4.4 static): one shared position for every row. Optional left padding
+  (`pad_lens`) lets prompts of different lengths share a batch: pad keys are masked.
+- SlotKVCache (4.4 continuous): S independent slots, each with its own length, so
+  requests can join and leave between decode steps.
+
+Buffers are [B, n_kv_heads, max_len, head_dim] per layer. No reference cycles: the
+caches own plain tensor lists and `layer(i)` returns a short-lived view. (A cycle once
+kept every finished cache's GPU memory alive until Python's cycle collector ran.)
+
+Left padding needs no per-row RoPE positions: RoPE makes q·k depend only on the
+distance between two positions, and padding shifts every real token of a row by the
+same amount. Pad *queries* may attend to themselves, so no row is fully masked (a fully
+masked softmax row is NaN, and NaN would leak through 0 × NaN in later layers).
+"""
+
+import torch
+from torch import Tensor
+
+
+def causal_kwargs(start: int, T: int, device) -> dict:
+    """SDPA masking for T new queries at positions start..start+T-1 over start+T keys, no padding.
+
+    - start 0 (prefill / training): standard causal.
+    - one new token: it may see every cached key, so no mask.
+    - several new tokens on a non-empty cache (chunked prefill, speculative verify):
+      SDPA's is_causal aligns to the top-left, which is wrong here, so pass an explicit
+      mask where query i sees keys 0..start+i.
+    """
+    if start == 0:
+        return {"is_causal": True}
+    if T == 1:
+        return {}
+    q_pos = torch.arange(start, start + T, device=device)[:, None]
+    k_pos = torch.arange(start + T, device=device)[None, :]
+    return {"attn_mask": k_pos <= q_pos}
+
+
+def padded_causal_mask(start: int, T: int, first_valid: Tensor) -> Tensor:
+    """(B, 1, T, start+T) bool: causal, keys before first_valid[b] hidden, diagonal always kept."""
+    dev = first_valid.device
+    q_pos = torch.arange(start, start + T, device=dev)[:, None]          # (T, 1)
+    k_pos = torch.arange(start + T, device=dev)[None, :]                 # (1, K)
+    causal = k_pos <= q_pos
+    valid = k_pos[None] >= first_valid[:, None, None]                    # (B, 1, K)
+    return (causal[None] & (valid | (k_pos == q_pos)[None]))[:, None]    # (B, 1, T, K)
+
+
+class KVCache:
+    """Shared-position cache. pad_lens (B,): left-pad count per row, or None for no padding."""
+
+    def __init__(self, n_layers: int, max_len: int, pad_lens: Tensor | None = None):
+        self.max_len = max_len
+        self.pos = 0
+        self.pad_lens = pad_lens
+        self.k: list[Tensor | None] = [None] * n_layers
+        self.v: list[Tensor | None] = [None] * n_layers
+        self.attn_kwargs: dict = {}
+
+    def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
+        cos, sin = rope(T, self.pos)
+        if self.pad_lens is None:
+            self.attn_kwargs = causal_kwargs(self.pos, T, cos.device)
+        else:
+            self.attn_kwargs = {"attn_mask": padded_causal_mask(self.pos, T, self.pad_lens)}
+        return cos, sin
+
+    def end_step(self, T: int) -> None:
+        self.pos += T
+
+    def layer(self, i: int) -> "LayerKV":
+        return LayerKV(self, i)
+
+    def nbytes(self) -> int:
+        return sum(t.numel() * t.element_size() for t in self.k + self.v if t is not None)
+
+
+class LayerKV:
+    """One layer's view of a KVCache."""
+
+    def __init__(self, cache: KVCache, layer_idx: int):
+        self.cache = cache
+        self.i = layer_idx
+
+    @property
+    def start(self) -> int:
+        """Position of the first new token = number of tokens already cached."""
+        return self.cache.pos
+
+    @property
+    def attn_kwargs(self) -> dict:
+        return self.cache.attn_kwargs
+
+    def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        """Store k, v (B, H, T, D) at positions start..start+T-1; return all cached K, V so far."""
+        c, i = self.cache, self.i
+        B, H, T, D = k.shape
+        if c.k[i] is None:
+            shape = (B, H, c.max_len, D)
+            c.k[i] = torch.zeros(shape, dtype=k.dtype, device=k.device)
+            c.v[i] = torch.zeros(shape, dtype=v.dtype, device=v.device)
+        start, end = self.start, self.start + T
+        if end > c.max_len:
+            raise ValueError(f"KV cache full: position {end} > max_len {c.max_len}")
+        c.k[i][:, :, start:end] = k
+        c.v[i][:, :, start:end] = v
+        return c.k[i][:, :, :end], c.v[i][:, :, :end]
+
+
+class SlotKVCache:
+    """S slots, each holding one sequence with its own length (continuous batching, 4.4).
+
+    Slot b holds keys at positions first[b] .. lens[b]-1 (positions before first[b] are
+    left padding from a batched prefill). A decode step feeds one token per slot, at
+    position lens[b]. Free slots still run (their output is ignored) so shapes stay fixed.
+
+    Every slot reserves max_len positions whatever its request needs: the waste that
+    paged KV (4.5) removes. `used_tokens()` vs `reserved_tokens()` measures it.
+    """
+
+    def __init__(self, n_layers: int, n_slots: int, n_kv_heads: int, head_dim: int,
+                 max_len: int, dtype: torch.dtype, device):
+        shape = (n_slots, n_kv_heads, max_len, head_dim)
+        self.k = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
+        self.v = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
+        self.n_slots, self.max_len = n_slots, max_len
+        # CPU copies drive scheduling without GPU syncs; device copies feed RoPE and masks.
+        self.first = [0] * n_slots
+        self.lens = [0] * n_slots
+        self.active = [False] * n_slots
+        self._lens_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._first_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._active_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._rows = torch.arange(n_slots, device=device)
+        self.attn_kwargs: dict = {}
+        self._end = 0
+        # 4.11b: a fixed attention window (a graph bucket) instead of "as far as the longest
+        # slot", so the step's shapes don't change between steps. None = eager default.
+        self.window: int | None = None
+
+    # -- scheduling --------------------------------------------------------
+    def free_slots(self) -> list[int]:
+        return [s for s in range(self.n_slots) if not self.active[s]]
+
+    def admit(self, slots: list[int], prefill: KVCache, pad_lens: list[int]) -> None:
+        """Copy a finished batched prefill (row j → slot slots[j]) into the slots."""
+        L = prefill.pos
+        if L > self.max_len:
+            raise ValueError(f"prompt length {L} > slot max_len {self.max_len}")
+        idx = torch.tensor(slots, device=self._lens_dev.device)
+        for i in range(len(self.k)):
+            self.k[i][idx, :, :L] = prefill.k[i][:, :, :L]
+            self.v[i][idx, :, :L] = prefill.v[i][:, :, :L]
+        for s, p in zip(slots, pad_lens):
+            self.first[s], self.lens[s], self.active[s] = p, L, True
+        self._sync()
+
+    def release(self, slot: int) -> None:
+        self.active[slot] = False
+        self.first[slot] = self.lens[slot] = 0
+        self._sync()
+
+    def _sync(self) -> None:
+        """Push the CPU bookkeeping to the device (only on admit / release, not every step)."""
+        self._lens_dev.copy_(torch.tensor(self.lens))
+        self._first_dev.copy_(torch.tensor(self.first))
+        self._active_dev.copy_(torch.tensor(self.active, dtype=torch.long))
+
+    def can_admit(self, lengths: list[int]) -> bool:
+        return True                                          # a free slot is all a request needs
+
+    def reserve_step(self) -> list[int]:
+        return []                                            # never preempts: capacity is per slot
+
+    def used_tokens(self) -> int:
+        return sum(n - f for n, f, a in zip(self.lens, self.first, self.active) if a)
+
+    def allocated_tokens(self) -> int:
+        """KV positions claimed by running sequences (a whole slot each)."""
+        return sum(self.active) * self.max_len
+
+    def reserved_tokens(self) -> int:
+        return self.n_slots * self.max_len
+
+    def nbytes(self) -> int:
+        return sum(t.numel() * t.element_size() for t in self.k + self.v)
+
+    # -- model protocol (decode only: one token per slot) ------------------
+    def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
+        if T != 1:
+            raise ValueError("SlotKVCache decodes one token per slot; prefill goes through KVCache")
+        if max(self.lens) >= self.max_len:
+            raise ValueError(f"a slot is full (max_len {self.max_len})")
+        need = max(self.lens) + 1                            # attend only as far as the longest slot
+        if self.window is not None and self.window < need:
+            raise ValueError(f"window {self.window} < {need} keys")
+        self._end = self.window or need
+        k_pos = torch.arange(self._end, device=self._lens_dev.device)
+        valid = (k_pos >= self._first_dev[:, None]) & (k_pos <= self._lens_dev[:, None])
+        self.attn_kwargs = {"attn_mask": valid[:, None, None, :]}   # (S, 1, 1, end)
+        return rope.at(self._lens_dev[:, None])
+
+    def end_step(self, T: int) -> None:
+        self.advance_host()
+        self._lens_dev += self._active_dev                   # same update on the device, no copy
+
+    def advance_host(self) -> None:
+        """The host half of end_step. A CUDA-graph replay runs only the device half, so the
+        graph decoder calls this after each replay."""
+        for s in range(self.n_slots):
+            if self.active[s]:
+                self.lens[s] += 1
+
+    def layer(self, i: int) -> "SlotLayerKV":
+        return SlotLayerKV(self, i)
+
+
+class SlotLayerKV:
+    def __init__(self, cache: SlotKVCache, layer_idx: int):
+        self.cache = cache
+        self.i = layer_idx
+
+    @property
+    def attn_kwargs(self) -> dict:
+        return self.cache.attn_kwargs
+
+    def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        """k, v (S, H, 1, D): write each slot's token at its own position lens[s]."""
+        c, i = self.cache, self.i
+        c.k[i][c._rows, :, c._lens_dev] = k[:, :, 0]
+        c.v[i][c._rows, :, c._lens_dev] = v[:, :, 0]
+        return c.k[i][:, :, :c._end], c.v[i][:, :, :c._end]

src/tinyllm/logger.py

@@ -0,0 +1,58 @@
+"""Experiment logging behind one small interface (D-013).
+
+The training loop only calls log / log_sample / finish. Switching backend (W&B now,
+TensorBoard as the fallback) means adding one class here. Only scalars and text
+samples are logged, never checkpoints (W&B Free plan: 5 GB/month).
+
+W&B login happens once, inside WSL: `uv run wandb login`. The key goes to ~/.netrc,
+never into the repo.
+"""
+
+
+class Logger:
+    """No-op logger (logger="none"): used by the tests and for quick local runs."""
+
+    run_id: str | None = None
+
+    def log(self, metrics: dict[str, float], step: int) -> None:
+        pass
+
+    def log_sample(self, text: str, step: int) -> None:
+        pass
+
+    def finish(self) -> None:
+        pass
+
+
+class WandbLogger(Logger):
+    def __init__(self, project: str, name: str, config: dict, mode: str = "online",
+                 run_id: str | None = None, dir: str | None = None):
+        import wandb  # imported lazily: the "none" logger never needs it
+
+        self._wandb = wandb
+        # Passing the stored run id on resume continues the same W&B run (same charts).
+        self.run = wandb.init(project=project, name=name, config=config, mode=mode,
+                              id=run_id, resume="allow" if run_id else None, dir=dir)
+        self.run_id = self.run.id
+        self._samples: list[list] = []
+
+    def log(self, metrics: dict[str, float], step: int) -> None:
+        self.run.log(metrics, step=step)
+
+    def log_sample(self, text: str, step: int) -> None:
+        # W&B tables are immutable, so re-log the growing table: one row per eval.
+        self._samples.append([step, text])
+        self.run.log({"samples": self._wandb.Table(columns=["step", "text"], data=self._samples)},
+                     step=step)
+
+    def finish(self) -> None:
+        self.run.finish()
+
+
+def make_logger(backend: str, **kwargs) -> Logger:
+    """kwargs are WandbLogger arguments; the no-op logger ignores them."""
+    if backend == "none":
+        return Logger()
+    if backend == "wandb":
+        return WandbLogger(**kwargs)
+    raise ValueError(f"unknown logger {backend!r}; use 'wandb' or 'none'")

src/tinyllm/logit_compare.py

@@ -0,0 +1,64 @@
+"""Compare next-token logits between models (step 4.9c; reused by 4.8b).
+
+Logits are collected the way inference produces them: a prefill over the first `n_prefill`
+tokens, then one token per step through the KV cache (teacher-forced, so every model sees
+the same sequence). Nothing here imports transformers: an HF model is passed in and only
+called as `hf(input_ids, past_key_values=..., use_cache=True)` (D-015: test reference only).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from tinyllm.kv_cache import KVCache
+
+
+@torch.no_grad()
+def ours_stepwise(model, ids: Tensor, n_prefill: int) -> Tensor:
+    """(B, T, V) fp32 CPU logits: prefill ids[:, :n_prefill], then one decode step per token."""
+    cache = KVCache(model.cfg.n_layers, model.cfg.max_seq_len)
+    autocast = dict(device_type=ids.device.type, dtype=torch.bfloat16, enabled=ids.device.type == "cuda")
+    out = []
+    with torch.autocast(**autocast):                      # same setting as generate()
+        out.append(model(ids[:, :n_prefill], cache=cache)[0].float().cpu())
+        for t in range(n_prefill, ids.shape[1]):
+            out.append(model(ids[:, t:t + 1], cache=cache)[0].float().cpu())
+    return torch.cat(out, dim=1)
+
+
+@torch.no_grad()
+def hf_stepwise(hf, ids: Tensor, n_prefill: int) -> Tensor:
+    """Same as ours_stepwise for an HF causal LM, using its own KV cache."""
+    o = hf(ids[:, :n_prefill], use_cache=True)
+    out, pkv = [o.logits.float().cpu()], o.past_key_values
+    for t in range(n_prefill, ids.shape[1]):
+        o = hf(ids[:, t:t + 1], past_key_values=pkv, use_cache=True)
+        out.append(o.logits.float().cpu())
+        pkv = o.past_key_values
+    return torch.cat(out, dim=1)
+
+
+def logit_stats(ref: Tensor, x: Tensor) -> dict[str, float]:
+    """x against a reference, over all positions: abs error, KL(ref || x) in nats, top-1 agreement."""
+    ref, x = ref.float().flatten(0, -2), x.float().flatten(0, -2)
+    d = (x - ref).abs()
+    lp_ref, lp_x = F.log_softmax(ref, -1), F.log_softmax(x, -1)
+    kl = (lp_ref.exp() * (lp_ref - lp_x)).sum(-1)
+    return {
+        "max_abs": d.max().item(),
+        "mean_abs": d.mean().item(),
+        "kl_mean": kl.mean().item(),
+        "kl_max": kl.max().item(),
+        "top1": (ref.argmax(-1) == x.argmax(-1)).float().mean().item(),
+        "positions": ref.shape[0],
+    }
+
+
+def first_divergence(a: list[int], b: list[int]) -> int | None:
+    """Index of the first differing token over the common length, or None if that prefix is equal."""
+    for i, (p, q) in enumerate(zip(a, b)):
+        if p != q:
+            return i
+    return None

src/tinyllm/logship.py

@@ -0,0 +1,220 @@
+"""Ship the server's logs to a private Hugging Face dataset (6.5b step 4, docs/SERVING.md "Privacy").
+
+The free Space wipes its disk on restart, so the request logs go to a private dataset:
+
+- `ShipHandler` writes a filtered copy of the `tinyllm` log as JSON lines, one file per server start
+  per UTC hour: `<dir>/<YYYY-MM-DD>/<start>-<replica>-<HH>.jsonl`. Files are never shared between
+  restarts, so an upload only ever touches the files of this run (in practice the current hour's).
+- Shipped: `request` and `slow request` lines, engine snapshots thinned to one per 30 s (the status
+  page keeps the 5 s ones in memory), the lifecycle lines (`model loaded`, `server ready`,
+  `shutting down: draining`, `stopped`) and errors without their traceback. Never debug lines.
+  No prompt or completion text and no raw IP reach the log at all (tests/test_logs.py checks it).
+- `LogShipper` uploads the folder with `huggingface_hub.CommitScheduler` every 10 minutes from a
+  background thread, plus once more during the graceful drain. Upload failures are logged to stderr
+  and never touch serving. A hard kill loses at most the last 10 minutes.
+- Retention: once a day (first at start) the uploader deletes dataset files older than 30 days, and
+  squashes the dataset history when its oldest commit is older than 30 days, so the deleted files
+  stop using storage.
+"""
+
+from __future__ import annotations
+
+import datetime as dt
+import logging
+import os
+import re
+import shutil
+import threading
+import time
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Callable
+
+from tinyllm.obs import JsonFormatter
+
+log = logging.getLogger("tinyllm.logship")
+
+SHIPPED = frozenset({"model loaded", "server ready", "shutting down: draining", "stopped",
+                     "request", "slow request"})
+SNAPSHOT = "engine"
+KEEP_DAYS = 30
+UPLOAD_EVERY_MIN = 10.0
+SNAPSHOT_EVERY_S = 30.0
+PATH_IN_REPO = "logs"
+_DAY = re.compile(r"\d{4}-\d{2}-\d{2}")
+
+
+def start_id(t: float) -> str:
+    """The server start in a file name, UTC: 20261107T131500Z."""
+    return time.strftime("%Y%m%dT%H%M%SZ", time.gmtime(t))
+
+
+def default_replica() -> str:
+    """The Space name on HF (`SPACE_ID` is owner/name), else "local"; safe for a file name."""
+    name = os.environ.get("SPACE_ID", "local").split("/")[-1]
+    return re.sub(r"[^A-Za-z0-9_.-]", "-", name) or "local"
+
+
+def log_path(root: Path, start: str, replica: str, t: float) -> Path:
+    g = time.gmtime(t)
+    return root / time.strftime("%Y-%m-%d", g) / f"{start}-{replica}-{g.tm_hour:02d}.jsonl"
+
+
+class ShipHandler(logging.Handler):
+    """Appends the shipped lines to the hourly file. `lock` guards each write against the uploader
+    reading the file at the same moment (CommitScheduler caps the upload at the size it saw)."""
+
+    def __init__(self, root: Path, start: str, replica: str, snapshot_every_s: float = SNAPSHOT_EVERY_S,
+                 lock: Callable[[], object] | None = None):
+        super().__init__(level=logging.INFO)
+        self.root, self.start, self.replica = Path(root), start, replica
+        self.snapshot_every_s = snapshot_every_s
+        self._io_lock = lock or nullcontext
+        self._last_snapshot: dict[str, float] = {}
+        self.setFormatter(JsonFormatter(with_exc=False))
+
+    def keep(self, record: logging.LogRecord) -> bool:
+        msg = record.getMessage()
+        if msg == SNAPSHOT and record.levelno < logging.WARNING:
+            variant = str((getattr(record, "fields", None) or {}).get("variant"))
+            last = self._last_snapshot.get(variant)
+            if last is not None and record.created - last < 0.95 * self.snapshot_every_s:
+                return False
+            self._last_snapshot[variant] = record.created
+            return True
+        return msg in SHIPPED or record.levelno >= logging.ERROR
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            if not self.keep(record):
+                return
+            line = self.format(record) + "\n"
+            path = log_path(self.root, self.start, self.replica, record.created)
+            with self._io_lock():
+                path.parent.mkdir(parents=True, exist_ok=True)
+                with path.open("a", encoding="utf-8") as f:
+                    f.write(line)
+        except Exception:
+            self.handleError(record)
+
+
+def _day(name: str) -> dt.date | None:
+    m = _DAY.fullmatch(name)
+    try:
+        return dt.date.fromisoformat(m.group()) if m else None
+    except ValueError:
+        return None
+
+
+def apply_retention(api, repo_id: str, root: Path | None, now: float, keep_days: int = KEEP_DAYS) -> dict:
+    """Delete log files older than `keep_days` (by their UTC date folder) from the dataset, in one
+    commit, and from `root`; squash the dataset history once its oldest commit is older than
+    `keep_days`. `api` is an `HfApi` (a fake in the tests). Returns what was done."""
+    from huggingface_hub import CommitOperationDelete
+
+    cutoff = dt.datetime.fromtimestamp(now, dt.timezone.utc).date() - dt.timedelta(days=keep_days)
+    old = []
+    for f in api.list_repo_files(repo_id, repo_type="dataset"):
+        parts = f.split("/")
+        if len(parts) >= 3 and parts[0] == PATH_IN_REPO and (d := _day(parts[1])) is not None and d < cutoff:
+            old.append(f)
+    if old:
+        api.create_commit(repo_id, repo_type="dataset", operations=[CommitOperationDelete(f) for f in old],
+                          commit_message=f"Delete logs older than {keep_days} days")
+    commits = api.list_repo_commits(repo_id, repo_type="dataset")
+    squash = len(commits) > 1 and commits[-1].created_at.timestamp() < now - keep_days * 86400
+    if squash:
+        api.super_squash_history(repo_id, repo_type="dataset",
+                                 commit_message=f"Squash history (logs kept {keep_days} days)")
+    local = 0
+    if root is not None and root.is_dir():
+        for d in root.iterdir():
+            if d.is_dir() and (day := _day(d.name)) is not None and day < cutoff:
+                shutil.rmtree(d, ignore_errors=True)
+                local += 1
+    return {"deleted": len(old), "squashed": squash, "local_dirs_deleted": local}
+
+
+def _scheduler_class():
+    """CommitScheduler that runs the retention once a day before an upload. Uploads, deletes and the
+    squash all run on the scheduler's one worker thread, so they never overlap."""
+    from huggingface_hub import CommitScheduler
+
+    class RetentionScheduler(CommitScheduler):
+        def __init__(self, *a, keep_days: int = KEEP_DAYS, clock: Callable[[], float] = time.time, **kw):
+            self.keep_days, self.clock, self.last_retention = keep_days, clock, None
+            super().__init__(*a, **kw)
+
+        def push_to_hub(self):
+            now = self.clock()
+            if self.last_retention is None or now - self.last_retention >= 86400:
+                self.last_retention = now
+                try:
+                    done = apply_retention(self.api, self.repo_id, self.folder_path, now, self.keep_days)
+                    log.info("log retention", extra={"fields": {"dataset": self.repo_id, **done}})
+                except Exception as e:
+                    log.warning("log retention failed", extra={"fields": {"error": f"{type(e).__name__}: {e}"}})
+            return super().push_to_hub()
+
+    return RetentionScheduler
+
+
+class LogShipper:
+    """Local hourly JSONL under `log_dir`; with `dataset`, uploads it to that private HF dataset.
+    The token is the usual one: `HF_TOKEN` (a Space secret) or the local `hf auth login`."""
+
+    def __init__(self, log_dir: str | Path, dataset: str | None = None, replica: str | None = None,
+                 every_min: float = UPLOAD_EVERY_MIN, keep_days: int = KEEP_DAYS,
+                 snapshot_every_s: float = SNAPSHOT_EVERY_S, start: float | None = None,
+                 scheduler_factory: Callable | None = None):
+        self.root = Path(log_dir)
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.dataset, self.every_min, self.keep_days = dataset, every_min, keep_days
+        self.start = start_id(time.time() if start is None else start)
+        self.replica = replica or default_replica()
+        self.scheduler = None
+        self._factory = scheduler_factory
+        self._lock = threading.Lock()
+        self._thread: threading.Thread | None = None
+        self.handler = ShipHandler(self.root, self.start, self.replica, snapshot_every_s, self._io_lock)
+
+    def _io_lock(self):
+        s = self.scheduler
+        return s.lock if s is not None else self._lock
+
+    def attach(self, logger: logging.Logger | None = None) -> "LogShipper":
+        (logger or logging.getLogger("tinyllm")).addHandler(self.handler)
+        return self
+
+    def start_upload(self) -> None:
+        """Open the dataset in the background (the first upload and the retention run right away)."""
+        if self.dataset and self._thread is None:
+            self._thread = threading.Thread(target=self._connect, name="log-upload", daemon=True)
+            self._thread.start()
+
+    def _connect(self) -> None:
+        factory = self._factory or _scheduler_class()
+        try:
+            self.scheduler = factory(repo_id=self.dataset, repo_type="dataset", private=True,
+                                     folder_path=self.root, path_in_repo=PATH_IN_REPO,
+                                     every=self.every_min, keep_days=self.keep_days)
+            log.info("log upload on", extra={"fields": {"dataset": self.dataset, "every_min": self.every_min}})
+        except Exception as e:      # no token, no access, no network: keep the logs local
+            log.warning("log upload off: logs stay local", extra={"fields": {
+                "dataset": self.dataset, "dir": str(self.root), "error": f"{type(e).__name__}: {e}"}})
+
+    def close(self, timeout_s: float = 30.0) -> None:
+        """Last upload (during the graceful drain), then stop the scheduler. Never raises."""
+        if self._thread is not None:
+            self._thread.join(timeout_s)
+        s = self.scheduler
+        if s is None:
+            return
+        try:
+            s.trigger().result(timeout=timeout_s)
+            log.info("logs uploaded", extra={"fields": {"dataset": self.dataset}})
+        except Exception as e:
+            log.warning("last log upload failed", extra={"fields": {"error": f"{type(e).__name__}: {e}"}})
+        finally:
+            s.stop()
+        logging.getLogger("tinyllm").removeHandler(self.handler)

src/tinyllm/model.py

@@ -0,0 +1,245 @@
+"""Llama-style decoder components, written from scratch (docs/ARCHITECTURE.md).
+
+Parameter names follow HF Llama (weight, gate_proj, up_proj, down_proj) so the
+Phase 4.8 export, and loading TinyLlama in Phase 4.9, need no renaming.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from tinyllm.config import ModelConfig
+
+
+class RMSNorm(nn.Module):
+    """y = x / sqrt(mean(x^2) + eps) * weight
+
+    Like LayerNorm, but with no mean subtraction and no bias: cheaper, and it
+    works just as well in LLMs. The statistic is computed in fp32 even under
+    bf16, because a sum of squares loses precision in bf16.
+
+    Backend "triton" (step 4.10a, the default since D-032) runs tinyllm.kernels.rmsnorm: one kernel
+    instead of about seven, and add_forward fuses the residual add before it. It is used only for
+    CUDA inputs when no gradient is needed (the kernel has no backward); otherwise, and on backend
+    "torch", this PyTorch code runs. So training and CPU inference are unchanged; GPU inference
+    differs from the PyTorch code only by bf16 rounding (docs/INFERENCE.md, 4.10a).
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-5, backend: str | None = None):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.backend = _check_norm_backend(DEFAULT_NORM_BACKEND if backend is None else backend)
+
+    def _kernel(self, x: Tensor):
+        if self.backend != "triton" or not x.is_cuda:
+            return None
+        if torch.is_grad_enabled() and (x.requires_grad or self.weight.requires_grad):
+            return None
+        return _norm_kernel()
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (k := self._kernel(x)) is not None:
+            return k.rmsnorm(x, self.weight, self.eps)
+        xf = x.float()
+        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        return xf.to(x.dtype) * self.weight
+
+    def add_forward(self, x: Tensor, a: Tensor) -> tuple[Tensor, Tensor]:
+        """(self(x + a), x + a): the residual add and the norm after it, fused on backend "triton"."""
+        if (k := self._kernel(x)) is not None and not (torch.is_grad_enabled() and a.requires_grad):
+            return k.add_rmsnorm(x, a, self.weight, self.eps)
+        h = x + a
+        return self(h), h
+
+    def extra_repr(self) -> str:
+        return f"{self.weight.shape[0]}, eps={self.eps}, backend={self.backend}"
+
+
+NORM_BACKENDS = ("torch", "triton")
+DEFAULT_NORM_BACKEND = "triton"   # D-032: the fused kernel by default; it steps aside for autograd and CPU
+
+
+def _check_norm_backend(backend: str) -> str:
+    if backend not in NORM_BACKENDS:
+        raise ValueError(f"unknown RMSNorm backend {backend!r}; one of {NORM_BACKENDS}")
+    return backend
+
+
+_NORM_KERNEL: list = []
+
+
+def _norm_kernel():
+    """tinyllm.kernels.rmsnorm, or None without Triton (imported once, on first CUDA use)."""
+    if not _NORM_KERNEL:
+        try:
+            from tinyllm.kernels import rmsnorm
+        except ImportError:
+            rmsnorm = None
+        _NORM_KERNEL.append(rmsnorm)
+    return _NORM_KERNEL[0]
+
+
+def set_norm_backend(model: nn.Module, backend: str) -> int:
+    """Switch every RMSNorm in `model` to `backend` (step 4.10a). Returns how many."""
+    _check_norm_backend(backend)
+    norms = [m for m in model.modules() if isinstance(m, RMSNorm)]
+    for m in norms:
+        m.backend = backend
+    return len(norms)
+
+
+def norm_backend(model: nn.Module | None = None) -> str:
+    """The RMSNorm code `model` runs at inference here: "triton" only if its norms ask for it, its
+    weights are on CUDA and Triton imports; else "torch" ("mixed" if its norms disagree). With no
+    model: what a freshly built model on CUDA would run. Recorded in benchmark provenance (D-032)."""
+    if model is None:
+        kinds = {DEFAULT_NORM_BACKEND}
+        on_cuda = torch.cuda.is_available()
+    else:
+        kinds = {m.backend for m in model.modules() if isinstance(m, RMSNorm)}
+        p = next(model.parameters(), None)
+        on_cuda = p is not None and p.is_cuda
+    if len(kinds) > 1:
+        return "mixed"
+    if kinds == {"triton"} and on_cuda and _norm_kernel() is not None:
+        return "triton"
+    return "torch"
+
+
+class RotaryEmbedding(nn.Module):
+    """Precomputed RoPE cos/sin tables.
+
+    Each pair of channels (i, i + head_dim/2) in q and k is rotated by the angle
+    pos * theta^(-2i/head_dim). Then q·k depends only on the *relative* distance
+    between positions, and no position parameters are learned.
+    This is the "rotate_half" layout that HF Llama uses, so TinyLlama weights
+    work unchanged.
+    """
+
+    def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10000.0):
+        super().__init__()
+        inv_freq = theta ** (-torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)  # (D/2,)
+        pos = torch.arange(max_seq_len, dtype=torch.float32)
+        angles = torch.outer(pos, inv_freq)              # (T, D/2)
+        angles = torch.cat([angles, angles], dim=-1)     # (T, D): same angle for i and i + D/2
+        # Not saved in checkpoints (persistent=False): they're derived from the config.
+        self.register_buffer("cos", angles.cos(), persistent=False)
+        self.register_buffer("sin", angles.sin(), persistent=False)
+
+    def forward(self, seq_len: int, offset: int = 0) -> tuple[Tensor, Tensor]:
+        """cos/sin for positions offset .. offset+seq_len-1. The offset is for KV-cached decode."""
+        end = offset + seq_len
+        if end > self.cos.shape[0]:
+            raise ValueError(f"position {end} exceeds max_seq_len {self.cos.shape[0]}")
+        return self.cos[offset:end], self.sin[offset:end]
+
+    def at(self, positions: Tensor) -> tuple[Tensor, Tensor]:
+        """cos/sin for per-row positions (B, T) → (B, 1, T, head_dim), broadcasting over heads.
+        Used when rows sit at different positions (continuous batching, 4.4).
+        No bounds check here (it would force a GPU sync every step): the caller checks."""
+        return self.cos[positions][:, None], self.sin[positions][:, None]
+
+
+def rotate_half(x: Tensor) -> Tensor:
+    """[x1, x2] -> [-x2, x1], with x1, x2 the two halves of the last dim."""
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+    """x: (B, n_heads, T, head_dim); cos/sin: (T, head_dim)."""
+    return x * cos.to(x.dtype) + rotate_half(x) * sin.to(x.dtype)
+
+
+class SwiGLU(nn.Module):
+    """down(silu(gate(x)) * up(x)): a gated MLP, used by Llama, Mistral, Qwen.
+
+    It has 3 matrices instead of 2, so ffn_hidden is ~(8/3)·d rather than 4·d
+    to keep the param count similar.
+    """
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
+        self.up_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
+        self.down_proj = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Block(nn.Module):
+    """Pre-norm decoder layer: x + attn(norm(x)), then x + mlp(norm(x))."""
+
+    def __init__(self, cfg: ModelConfig, layer_idx: int):
+        super().__init__()
+        # Imported here, not at the top: the attention modules import apply_rope from this file.
+        from tinyllm.attention import build_attention
+
+        self.input_layernorm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.self_attn = build_attention(cfg, layer_idx)
+        self.post_attention_layernorm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+
+    def forward(self, x: Tensor, rope: tuple[Tensor, Tensor], cache=None) -> Tensor:
+        a = self.self_attn(self.input_layernorm(x), rope, cache)
+        y, x = self.post_attention_layernorm.add_forward(x, a)  # x + a, then its norm (one kernel, 4.10a)
+        return x + self.mlp(y)
+
+
+# Projections that write into the residual stream. Their init is scaled down so the
+# stream's variance doesn't grow with depth (GPT-2 / nanoGPT recipe).
+RESIDUAL_PROJ = ("o_proj.weight", "down_proj.weight")
+
+
+class Transformer(nn.Module):
+    """Token ids (B, T) -> logits (B, T, vocab), plus the next-token loss when targets are given."""
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.layers = nn.ModuleList(Block(cfg, i) for i in range(cfg.n_layers))
+        self.norm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            # One matrix for both: saves V·d params (4.2M in main, 14% of the model).
+            self.lm_head.weight = self.embed_tokens.weight
+        self.rope = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        std = self.cfg.init_std
+        for m in self.modules():
+            if isinstance(m, (nn.Linear, nn.Embedding)):
+                nn.init.normal_(m.weight, mean=0.0, std=std)
+        resid_std = std / math.sqrt(2 * self.cfg.n_layers)  # 2 residual writes per layer
+        for name, p in self.named_parameters():
+            if name.endswith(RESIDUAL_PROJ):
+                nn.init.normal_(p, mean=0.0, std=resid_std)
+        # RMSNorm weights stay at 1.
+
+    def num_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())  # a tied matrix is counted once
+
+    def forward(
+        self, idx: Tensor, targets: Tensor | None = None, cache=None,
+    ) -> tuple[Tensor, Tensor | None]:
+        """cache: a tinyllm.kv_cache cache for incremental decoding (None in training).
+        With a cache, idx holds only the new tokens; the cache supplies their positions."""
+        _, T = idx.shape
+        rope = self.rope(T) if cache is None else cache.begin_step(T, self.rope)
+        x = self.embed_tokens(idx)
+        for i, layer in enumerate(self.layers):
+            x = layer(x, rope, None if cache is None else cache.layer(i))
+        if cache is not None:
+            cache.end_step(T)
+        logits = self.lm_head(self.norm(x))
+        loss = None
+        if targets is not None:
+            # Loss in fp32: log-softmax over 8192 classes is imprecise in bf16.
+            loss = F.cross_entropy(logits.float().flatten(0, 1), targets.flatten())
+        return logits, loss

src/tinyllm/obs.py

@@ -0,0 +1,362 @@
+"""Step 4.13 step 2: observability for the server (D-028, docs/SERVING.md).
+
+Standard library only, so the server and the Space need nothing extra.
+
+- RequestRecord: one JSON log line per finished request (timings, token counts, status).
+  Never the prompt or completion text; the client is a salted hash (hash_client).
+- EngineSnapshot: the engine's load, sampled every few seconds.
+- Metrics: thread-safe aggregator. Rolling windows (5 min, 1 h) back /stats (JSON);
+  cumulative counters and histograms back /metrics (Prometheus text format, so a real
+  Prometheus could scrape it unchanged). history() backs /stats/history for the status page
+  (6.5b): the last hour per minute, engine snapshots, and the SLOs with their error budgets.
+- JsonFormatter / setup_logging: every log line is one JSON object.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import math
+import sys
+import threading
+import time
+from collections import deque
+from dataclasses import asdict, dataclass, field
+
+WINDOWS = {"5m": 300.0, "1h": 3600.0}
+TTFT_BUCKETS = (0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0)
+LATENCY_BUCKETS = (0.1, 0.25, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0)
+
+# SLOs (docs/SERVING.md, 4.13 step 3). Targets are shares of good requests; the TTFT threshold
+# is per server (--slo-ttft-s: 1 s local GPU, 2 s on the Space).
+SLO_TARGETS = {"ttft": 0.95, "availability": 0.99, "shed": 0.99}
+SLO_WINDOW_S = 7 * 24 * 3600.0     # rolling 7 days, but only since the last restart (in memory)
+FAST_BURN = 14.4                   # Google SRE fast-burn threshold, judged on the last hour
+BURN_MIN_REQUESTS = 10             # fewer counted requests in the hour: no flag (too few to judge)
+HISTORY_MINUTES = 60
+
+
+def slo_counted(status: int) -> bool:
+    """Requests the SLOs judge: served (200), shed (429) and server failures (>= 500, incl. 504).
+    Client errors (400) and client disconnects (499) count for none of them."""
+    return status == 200 or status == 429 or status >= 500
+
+
+class _SloBucket:
+    """Counters for one hour of counted requests."""
+    __slots__ = ("counted", "ok", "slow", "failed", "shed")
+
+    def __init__(self):
+        self.counted = self.ok = self.slow = self.failed = self.shed = 0
+
+    def add(self, rec: "RequestRecord", ttft_s: float) -> None:
+        self.counted += 1
+        if rec.status == 200:
+            self.ok += 1
+            self.slow += rec.ttft_s is None or rec.ttft_s > ttft_s
+        elif rec.status == 429:
+            self.shed += 1
+        else:
+            self.failed += 1
+
+    def bad_and_total(self, slo: str) -> tuple[int, int]:
+        if slo == "ttft":
+            return self.slow, self.ok
+        return (self.failed, self.counted) if slo == "availability" else (self.shed, self.counted)
+
+
+def hash_client(ip: str, salt: str) -> str:
+    """Stable, salted id for a client. The raw IP is never stored."""
+    return hashlib.sha256(f"{salt}:{ip}".encode()).hexdigest()[:16]
+
+
+def percentile(xs, q: float) -> float | None:
+    """Linear-interpolated percentile (numpy's default); None for no data. q in [0, 100]."""
+    xs = sorted(xs)
+    if not xs:
+        return None
+    k = (len(xs) - 1) * q / 100
+    lo, hi = math.floor(k), math.ceil(k)
+    return xs[lo] + (xs[hi] - xs[lo]) * (k - lo)
+
+
+def itl_stats(token_times: list[float]) -> tuple[float | None, float | None]:
+    """Inter-token latency p50 / p95 from the times each token was produced."""
+    gaps = [b - a for a, b in zip(token_times, token_times[1:])]
+    return percentile(gaps, 50), percentile(gaps, 95)
+
+
+@dataclass
+class RequestRecord:
+    """One finished request (any outcome). Times in seconds; None when not reached."""
+    request_id: str
+    status: int                              # HTTP status: 200, 400, 429, 500, 504, 499 (client gone)
+    variant: str = "main"
+    model_version: str = ""
+    client: str = ""                         # hash_client(), never the IP
+    arrival_unix: float = 0.0
+    queue_wait_s: float | None = None        # arrival → admitted into a slot (first token)
+    ttft_s: float | None = None
+    itl_p50_s: float | None = None
+    itl_p95_s: float | None = None
+    latency_s: float | None = None
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    finish_reason: str | None = None         # "stop" (EOS), "length", "abort", "timeout", "error"
+    stream: bool = False
+    error: str | None = None
+    queue_depth_at_arrival: int = 0
+    busy_slots_at_arrival: int = 0
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+@dataclass
+class EngineSnapshot:
+    t_unix: float
+    queue_len: int
+    busy_slots: int
+    n_slots: int
+    tokens_per_s: float                      # generated tokens / s since the previous snapshot
+    kv_used_frac: float
+    ram_mb: float | None = None
+    vram_mb: float | None = None
+
+
+class _Histogram:
+    def __init__(self, buckets):
+        self.buckets = tuple(buckets)
+        self.counts = [0] * len(self.buckets)    # per bucket, not cumulative
+        self.sum = 0.0
+        self.count = 0
+
+    def observe(self, x: float) -> None:
+        for i, b in enumerate(self.buckets):
+            if x <= b:
+                self.counts[i] += 1
+                break
+        self.sum += x
+        self.count += 1
+
+    def lines(self, name: str, labels: str) -> list[str]:
+        sep = "," if labels else ""
+        out, cum = [], 0
+        for b, c in zip(self.buckets, self.counts):
+            cum += c
+            out.append(f'{name}_bucket{{{labels}{sep}le="{b:g}"}} {cum}')
+        out.append(f'{name}_bucket{{{labels}{sep}le="+Inf"}} {self.count}')
+        lab = f"{{{labels}}}" if labels else ""
+        out.append(f"{name}_sum{lab} {self.sum:.6f}")
+        out.append(f"{name}_count{lab} {self.count}")
+        return out
+
+
+@dataclass
+class _Window:
+    span_s: float
+    items: deque = field(default_factory=deque)   # (t_monotonic, RequestRecord)
+
+    def prune(self, now: float) -> None:
+        while self.items and self.items[0][0] < now - self.span_s:
+            self.items.popleft()
+
+
+class Metrics:
+    """Thread-safe: handlers call observe(), the engine thread calls snapshot()."""
+
+    def __init__(self, windows: dict[str, float] = WINDOWS, clock=time.monotonic, slo_ttft_s: float = 1.0,
+                 wall=time.time):
+        self._lock = threading.Lock()
+        self._clock, self._wall = clock, wall
+        self.slo_ttft_s = slo_ttft_s
+        self._windows = {k: _Window(v) for k, v in windows.items()}
+        self._hour = _Window(HISTORY_MINUTES * 60.0)        # per-minute history and the 1 h burn rate
+        self._slo: deque = deque()                          # (hour index since start, _SloBucket)
+        self._snaps: deque = deque()                        # (t_monotonic, EngineSnapshot), last hour
+        self._requests: dict[tuple[str, int], int] = {}      # (variant, status) → count
+        self._tokens: dict[str, int] = {}                     # variant → completion tokens
+        self._ttft: dict[str, _Histogram] = {}
+        self._latency: dict[str, _Histogram] = {}
+        self.last_snapshot: EngineSnapshot | None = None
+        self.started = clock()
+
+    def observe(self, rec: RequestRecord, now: float | None = None) -> None:
+        now = self._clock() if now is None else now
+        with self._lock:
+            for w in (*self._windows.values(), self._hour):
+                w.items.append((now, rec))
+                w.prune(now)
+            if slo_counted(rec.status):
+                hour = int((now - self.started) // 3600)
+                if not self._slo or self._slo[-1][0] != hour:
+                    self._slo.append((hour, _SloBucket()))
+                    while self._slo and self._slo[0][0] <= hour - SLO_WINDOW_S / 3600:
+                        self._slo.popleft()
+                self._slo[-1][1].add(rec, self.slo_ttft_s)
+            v = rec.variant
+            self._requests[(v, rec.status)] = self._requests.get((v, rec.status), 0) + 1
+            self._tokens[v] = self._tokens.get(v, 0) + rec.completion_tokens
+            if rec.ttft_s is not None:
+                self._ttft.setdefault(v, _Histogram(TTFT_BUCKETS)).observe(rec.ttft_s)
+            if rec.status == 200 and rec.latency_s is not None:
+                self._latency.setdefault(v, _Histogram(LATENCY_BUCKETS)).observe(rec.latency_s)
+
+    def snapshot(self, snap: EngineSnapshot, now: float | None = None) -> None:
+        now = self._clock() if now is None else now
+        with self._lock:
+            self.last_snapshot = snap
+            self._snaps.append((now, snap))
+            while self._snaps and self._snaps[0][0] < now - HISTORY_MINUTES * 60:
+                self._snaps.popleft()
+
+    def history(self, now: float | None = None) -> dict:
+        """/stats/history, for the status page: the last hour per minute, the engine snapshots of the
+        last hour, and each SLO over the rolling 7-day window (since the last restart)."""
+        now = self._clock() if now is None else now
+        wall = self._wall()
+        with self._lock:
+            self._hour.prune(now)
+            recs = list(self._hour.items)
+            snaps = [(t, s) for t, s in self._snaps if t >= now - HISTORY_MINUTES * 60]
+            hour_now = int((now - self.started) // 3600)
+            buckets = [b for h, b in self._slo if h > hour_now - SLO_WINDOW_S / 3600]
+        uptime = now - self.started
+
+        by_min: list[list[RequestRecord]] = [[] for _ in range(HISTORY_MINUTES)]
+        for t, r in recs:
+            ago = int((now - t) // 60)
+            if 0 <= ago < HISTORY_MINUTES:
+                by_min[ago].append(r)
+        minutes = []
+        for ago in range(HISTORY_MINUTES - 1, -1, -1):        # oldest first
+            rs = by_min[ago]
+            ok = [r for r in rs if r.status == 200]
+            ttft = [r.ttft_s for r in ok if r.ttft_s is not None]
+            lat = [r.latency_s for r in ok if r.latency_s is not None]
+            minutes.append({
+                "minutes_ago": ago, "covered": uptime >= ago * 60,   # the server was up for this minute
+                "requests": len(rs), "ok": len(ok),
+                "failed": sum(r.status >= 500 for r in rs), "shed": sum(r.status == 429 for r in rs),
+                "completion_tokens": sum(r.completion_tokens for r in rs),
+                "ttft_p50_s": percentile(ttft, 50), "ttft_p95_s": percentile(ttft, 95),
+                "latency_p95_s": percentile(lat, 95)})
+
+        hour = _SloBucket()
+        for _, r in recs:
+            if slo_counted(r.status):
+                hour.add(r, self.slo_ttft_s)
+        slos = {}
+        for name, target in SLO_TARGETS.items():
+            bad, total = map(sum, zip(*(b.bad_and_total(name) for b in buckets))) if buckets else (0, 0)
+            hbad, htotal = hour.bad_and_total(name)
+            allowed = 1 - target
+            burn = (hbad / htotal) / allowed if htotal else None
+            slos[name] = {
+                "target": target, "good": total - bad, "total": total,
+                "compliance": round((total - bad) / total, 6) if total else None,
+                "met": (total - bad) / total >= target if total else None,
+                "budget_left": round(1 - bad / (allowed * total), 6) if total else None,
+                "burn_rate_1h": round(burn, 3) if burn is not None else None,
+                "requests_1h": htotal,
+                "fast_burn": burn is not None and htotal >= BURN_MIN_REQUESTS and burn > FAST_BURN}
+        slos["ttft"]["threshold_s"] = self.slo_ttft_s
+        slos["shed"]["judged"] = False          # only at or below the rated load, unknown for the Space yet
+
+        return {"now_unix": round(wall, 3), "uptime_s": round(uptime, 1),
+                "slo_window_s": round(min(SLO_WINDOW_S, uptime), 1), "minutes": minutes,
+                "snapshots": [{"seconds_ago": round(now - t, 1), "queue_len": s.queue_len,
+                               "busy_slots": s.busy_slots, "n_slots": s.n_slots,
+                               "tokens_per_s": round(s.tokens_per_s, 1)} for t, s in snaps],
+                "slo": slos, "fast_burn_threshold": FAST_BURN, "burn_min_requests": BURN_MIN_REQUESTS}
+
+    def stats(self, now: float | None = None) -> dict:
+        """/stats: per rolling window, request counts, rates and latency percentiles."""
+        now = self._clock() if now is None else now
+        out: dict = {"uptime_s": round(now - self.started, 1), "windows": {}}
+        with self._lock:
+            for name, w in self._windows.items():
+                w.prune(now)
+                recs = [r for _, r in w.items]
+                span = min(w.span_s, max(now - self.started, 1e-9))
+                ok = [r for r in recs if r.status == 200]
+                ttft = [r.ttft_s for r in ok if r.ttft_s is not None]
+                lat = [r.latency_s for r in ok if r.latency_s is not None]
+                qw = [r.queue_wait_s for r in ok if r.queue_wait_s is not None]
+                n = len(recs)
+                out["windows"][name] = {
+                    "requests": n,
+                    "requests_per_min": round(60 * n / span, 3),
+                    "ok": len(ok),
+                    "error_rate": round(sum(r.status >= 500 for r in recs) / n, 4) if n else 0.0,
+                    "reject_rate": round(sum(r.status == 429 for r in recs) / n, 4) if n else 0.0,
+                    "clients": len({r.client for r in recs}),
+                    "completion_tokens": sum(r.completion_tokens for r in recs),
+                    "ttft_s": {f"p{q}": percentile(ttft, q) for q in (50, 95, 99)},
+                    "latency_s": {f"p{q}": percentile(lat, q) for q in (50, 95, 99)},
+                    "queue_wait_s": {f"p{q}": percentile(qw, q) for q in (50, 95)},
+                }
+            out["engine"] = asdict(self.last_snapshot) if self.last_snapshot else None
+        return out
+
+    def prometheus(self) -> str:
+        """/metrics in the Prometheus text exposition format (version 0.0.4)."""
+        L: list[str] = []
+        with self._lock:
+            L += ["# HELP tinyllm_requests_total Finished requests by variant and HTTP status.",
+                  "# TYPE tinyllm_requests_total counter"]
+            for (v, s), c in sorted(self._requests.items()):
+                L.append(f'tinyllm_requests_total{{variant="{v}",status="{s}"}} {c}')
+            L += ["# HELP tinyllm_completion_tokens_total Generated tokens by variant.",
+                  "# TYPE tinyllm_completion_tokens_total counter"]
+            for v, c in sorted(self._tokens.items()):
+                L.append(f'tinyllm_completion_tokens_total{{variant="{v}"}} {c}')
+            for name, hists, help_ in (
+                    ("tinyllm_ttft_seconds", self._ttft, "Arrival to first token."),
+                    ("tinyllm_request_latency_seconds", self._latency, "Arrival to last token, successful requests.")):
+                L += [f"# HELP {name} {help_}", f"# TYPE {name} histogram"]
+                for v, h in sorted(hists.items()):
+                    L += h.lines(name, f'variant="{v}"')
+            s = self.last_snapshot
+            for name, help_, val in (
+                    ("tinyllm_queue_length", "Requests waiting for a slot.", s and s.queue_len),
+                    ("tinyllm_busy_slots", "Slots decoding a request.", s and s.busy_slots),
+                    ("tinyllm_slots", "Total slots.", s and s.n_slots),
+                    ("tinyllm_tokens_per_second", "Generated tokens per second (last snapshot).", s and s.tokens_per_s),
+                    ("tinyllm_kv_used_ratio", "Share of the KV cache holding tokens.", s and s.kv_used_frac)):
+                if s is not None:
+                    L += [f"# HELP {name} {help_}", f"# TYPE {name} gauge", f"{name} {val:g}"]
+        return "\n".join(L) + "\n"
+
+
+class JsonFormatter(logging.Formatter):
+    """One JSON object per line. Structured fields go in `extra={"fields": {...}}`.
+    `with_exc=False` leaves tracebacks out (the shipped logs, logship.py)."""
+
+    def __init__(self, with_exc: bool = True):
+        super().__init__()
+        self.with_exc = with_exc
+
+    def format(self, record: logging.LogRecord) -> str:
+        d = {"ts": round(record.created, 3), "level": record.levelname.lower(),
+             "logger": record.name, "msg": record.getMessage()}
+        d.update(getattr(record, "fields", None) or {})
+        if record.exc_info and self.with_exc:
+            d["exc"] = self.formatException(record.exc_info)
+        return json.dumps(d, default=str)
+
+
+def setup_logging(level: str = "info", path: str | None = None) -> logging.Logger:
+    """The `tinyllm` logger: JSON lines to stderr, and to `path` (JSONL) if given."""
+    log = logging.getLogger("tinyllm")
+    log.setLevel(level.upper())
+    log.handlers.clear()
+    log.propagate = False
+    handlers: list[logging.Handler] = [logging.StreamHandler(sys.stderr)]
+    if path:
+        handlers.append(logging.FileHandler(path, encoding="utf-8"))
+    for h in handlers:
+        h.setFormatter(JsonFormatter())
+        log.addHandler(h)
+    return log

src/tinyllm/paged_kv.py

@@ -0,0 +1,236 @@
+"""Step 4.5: paged KV cache (vLLM-style blocks), a drop-in for SlotKVCache in run_continuous.
+
+The analogy is virtual memory:
+    block table  = page table      (one per sequence: logical block j -> physical block)
+    block        = page            (block_size tokens of K and V, every layer)
+    token t      = (table[t // block_size], t % block_size)
+
+All layers share one block table: block b holds the same token positions in every layer's
+pool. Blocks are allocated as a sequence grows, not reserved up front, and every block has
+the same size, so a freed block fits any sequence: no external fragmentation, and at most
+block_size - 1 wasted slots per sequence (its last, partly filled block).
+
+Left padding from the batched prefill is not stored: only real tokens go into blocks.
+The pad count becomes a RoPE offset (`first`), so a token at storage index j keeps its
+rotation position first + j and the next query lands at first + len, as in SlotKVCache.
+
+Attention reads the blocks by gathering each sequence's blocks into a contiguous
+(S, H, n_blocks * block_size, D) view and calling SDPA as before. vLLM's PagedAttention
+kernel reads the blocks in place instead; the gather is this toy's extra copy (measured
+in docs/INFERENCE.md).
+
+When decode needs more blocks than are free, the newest sequence is preempted: its blocks
+are freed and the engine requeues it to be recomputed (prompt + tokens so far), which
+reproduces the same greedy tokens. This is vLLM's default recompute policy.
+"""
+
+import math
+
+import torch
+from torch import Tensor
+
+from tinyllm.kv_cache import KVCache
+
+
+class OutOfBlocks(RuntimeError):
+    pass
+
+
+class BlockAllocator:
+    """Fixed-size blocks with a free list. Lowest ids are handed out first (reused LIFO)."""
+
+    def __init__(self, n_blocks: int, block_size: int):
+        self.n_blocks, self.block_size = n_blocks, block_size
+        self._free = list(range(n_blocks - 1, -1, -1))     # stack: pop() gives block 0 first
+        self._in_use = [False] * n_blocks
+        self.peak_used = 0
+
+    @property
+    def num_free(self) -> int:
+        return len(self._free)
+
+    @property
+    def num_used(self) -> int:
+        return self.n_blocks - len(self._free)
+
+    def alloc(self) -> int:
+        if not self._free:
+            raise OutOfBlocks(f"all {self.n_blocks} KV blocks in use")
+        b = self._free.pop()
+        self._in_use[b] = True
+        self.peak_used = max(self.peak_used, self.num_used)
+        return b
+
+    def free(self, blocks: list[int]) -> None:
+        for b in blocks:
+            if not self._in_use[b]:
+                raise ValueError(f"block {b} freed twice (or never allocated)")
+            self._in_use[b] = False
+            self._free.append(b)
+
+
+class PagedKVCache:
+    """S sequence slots over a shared pool of n_blocks KV blocks.
+
+    Same model protocol as SlotKVCache (begin_step / layer / end_step, one token per slot)
+    and the same scheduling calls (free_slots / admit / release), plus can_admit() and
+    reserve_step() so the engine can respect the block budget.
+    """
+
+    def __init__(self, n_layers: int, n_slots: int, n_kv_heads: int, head_dim: int,
+                 n_blocks: int, block_size: int, max_len: int, dtype: torch.dtype, device):
+        self.alloc = BlockAllocator(n_blocks, block_size)
+        self.bs, self.n_slots, self.max_len = block_size, n_slots, max_len
+        self.max_blocks = math.ceil(max_len / block_size)
+        self.scratch = n_blocks           # extra block, never allocated: free slots write here
+        shape = (n_blocks + 1, n_kv_heads, block_size, head_dim)
+        self.k = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
+        self.v = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
+        # CPU bookkeeping (scheduling, no GPU syncs) and its device mirror (indexing, masks).
+        self.tables: list[list[int]] = [[] for _ in range(n_slots)]
+        self.lens = [0] * n_slots         # tokens stored (no padding)
+        self.first = [0] * n_slots        # RoPE offset = left-pad count at prefill
+        self.active = [False] * n_slots
+        self._admit_no = [0] * n_slots    # admission order, for picking preemption victims
+        self._admits = 0
+        self._table_dev = torch.full((n_slots, self.max_blocks), self.scratch,
+                                     dtype=torch.long, device=device)
+        self._lens_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._first_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._active_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
+        self._rows = torch.arange(n_slots, device=device)
+        self.attn_kwargs: dict = {}
+        self._n_read_blocks = 0
+        self._blk = self._off = None
+
+    # -- scheduling --------------------------------------------------------
+    def free_slots(self) -> list[int]:
+        return [s for s in range(self.n_slots) if not self.active[s]]
+
+    def can_admit(self, lengths: list[int]) -> bool:
+        """Enough free blocks for these prompts, plus room for each one's first decode token."""
+        return sum(math.ceil((n + 1) / self.bs) for n in lengths) <= self.alloc.num_free
+
+    def admit(self, slots: list[int], prefill: KVCache, pad_lens: list[int]) -> None:
+        """Copy row j of a padded prefill (real tokens only) into fresh blocks for slots[j]."""
+        L = prefill.pos
+        dev = self._lens_dev.device
+        rows, src, blk, off = [], [], [], []
+        for j, (s, p) in enumerate(zip(slots, pad_lens)):
+            n = L - p
+            if n >= self.max_len:
+                raise ValueError(f"prompt length {n} >= max_len {self.max_len}")
+            blocks = [self.alloc.alloc() for _ in range(math.ceil(n / self.bs))]
+            self.tables[s] = blocks
+            self.lens[s], self.first[s], self.active[s] = n, p, True
+            self._admits += 1
+            self._admit_no[s] = self._admits
+            self._table_dev[s].fill_(self.scratch)
+            self._table_dev[s, :len(blocks)] = torch.tensor(blocks, device=dev)
+            t = torch.arange(n)
+            rows.append(torch.full((n,), j))
+            src.append(t + p)
+            blk.append(torch.tensor(blocks, dtype=torch.long)[t // self.bs])
+            off.append(t % self.bs)
+        rows, src, blk, off = (torch.cat(x).to(dev) for x in (rows, src, blk, off))
+        for i in range(len(self.k)):
+            self.k[i][blk, :, off] = prefill.k[i][rows, :, src]     # (N, H, D) both sides
+            self.v[i][blk, :, off] = prefill.v[i][rows, :, src]
+        self._sync()
+
+    def release(self, slot: int) -> None:
+        self.alloc.free(self.tables[slot])
+        self.tables[slot] = []
+        self.active[slot] = False
+        self.lens[slot] = self.first[slot] = 0
+        self._table_dev[slot].fill_(self.scratch)
+        self._sync()
+
+    def reserve_step(self) -> list[int]:
+        """Make sure the next decode step's new blocks can be allocated. Preempts the newest
+        sequences (frees their blocks) until they can; returns the preempted slots, newest first."""
+        victims = []
+        while True:
+            need = sum(1 for s in range(self.n_slots)
+                       if self.active[s] and len(self.tables[s]) * self.bs <= self.lens[s])
+            if need <= self.alloc.num_free or not any(self.active):
+                return victims
+            v = max((s for s in range(self.n_slots) if self.active[s]), key=self._admit_no.__getitem__)
+            self.release(v)
+            victims.append(v)
+
+    def _sync(self) -> None:
+        self._lens_dev.copy_(torch.tensor(self.lens))
+        self._first_dev.copy_(torch.tensor(self.first))
+        self._active_dev.copy_(torch.tensor(self.active, dtype=torch.long))
+
+    def used_tokens(self) -> int:
+        return sum(n for n, a in zip(self.lens, self.active) if a)
+
+    def allocated_tokens(self) -> int:
+        return self.alloc.num_used * self.bs
+
+    def reserved_tokens(self) -> int:
+        """Pool capacity: the KV memory this cache holds."""
+        return self.alloc.n_blocks * self.bs
+
+    def nbytes(self) -> int:
+        return sum(t.numel() * t.element_size() for t in self.k + self.v)
+
+    # -- model protocol (decode only: one token per slot) ------------------
+    def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
+        if T != 1:
+            raise ValueError("PagedKVCache decodes one token per slot; prefill goes through KVCache")
+        if max(self.lens) >= self.max_len:
+            raise ValueError(f"a sequence is full (max_len {self.max_len})")
+        new_rows, new_cols, new_blks = [], [], []
+        for s in range(self.n_slots):                         # grow sequences that filled a block
+            if self.active[s] and len(self.tables[s]) * self.bs <= self.lens[s]:
+                b = self.alloc.alloc()                        # reserve_step made room
+                new_rows.append(s)
+                new_cols.append(len(self.tables[s]))
+                new_blks.append(b)
+                self.tables[s].append(b)
+        if new_rows:
+            dev = self._table_dev.device
+            self._table_dev[torch.tensor(new_rows, device=dev),
+                            torch.tensor(new_cols, device=dev)] = torch.tensor(new_blks, device=dev)
+        self._n_read_blocks = math.ceil((max(self.lens) + 1) / self.bs)
+        self._blk = self._table_dev[self._rows, self._lens_dev // self.bs]   # where to write
+        self._off = self._lens_dev % self.bs
+        k_pos = torch.arange(self._n_read_blocks * self.bs, device=self._lens_dev.device)
+        self.attn_kwargs = {"attn_mask": (k_pos <= self._lens_dev[:, None])[:, None, None, :]}
+        return rope.at((self._first_dev + self._lens_dev)[:, None])
+
+    def end_step(self, T: int) -> None:
+        for s in range(self.n_slots):
+            if self.active[s]:
+                self.lens[s] += 1
+        self._lens_dev += self._active_dev
+
+    def layer(self, i: int) -> "PagedLayerKV":
+        return PagedLayerKV(self, i)
+
+
+class PagedLayerKV:
+    def __init__(self, cache: PagedKVCache, layer_idx: int):
+        self.cache = cache
+        self.i = layer_idx
+
+    @property
+    def attn_kwargs(self) -> dict:
+        return self.cache.attn_kwargs
+
+    def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        """k, v (S, H, 1, D): write each slot's token into its block, then gather its blocks."""
+        c, i = self.cache, self.i
+        c.k[i][c._blk, :, c._off] = k[:, :, 0]
+        c.v[i][c._blk, :, c._off] = v[:, :, 0]
+        table = c._table_dev[:, :c._n_read_blocks]                     # (S, nb)
+        S, nb = table.shape
+        _, H, bs, D = c.k[i].shape
+
+        def gather(pool: Tensor) -> Tensor:                            # (S, nb, H, bs, D) -> (S, H, nb*bs, D)
+            return pool[table].transpose(1, 2).reshape(S, H, nb * bs, D)
+
+        return gather(c.k[i]), gather(c.v[i])

src/tinyllm/plots.py

@@ -0,0 +1,294 @@
+"""Step 6.2 (D-030): the README and docs/RESULTS.md plots, drawn from the same run as the tables.
+
+render_plots(data, curve) returns {file name: SVG text}; scripts/bench_all.py writes them to docs/img/.
+`data` is results.load_run(); `curve` is docs/data/train_curve.json (scripts/extract_train_curve.py).
+
+  - One transparent SVG per plot, readable on GitHub's light and dark themes: mid-gray text, and a
+    blue / orange pair checked for colour-vision deficiency on both backgrounds. Blue = ours, orange =
+    HF Transformers, gray = context. Every bar and line end carries its value, so color is never the
+    only key.
+  - Byte-identical between renders (fixed SVG hash salt, no date, text drawn as paths with the font
+    matplotlib ships), so re-rendering leaves git clean. matplotlib's version is pinned by uv.lock.
+  - Pyplot is not used: no global figure state, no GUI backend.
+"""
+
+from __future__ import annotations
+
+import io
+
+import matplotlib
+from matplotlib.figure import Figure
+from matplotlib.ticker import FixedLocator, FuncFormatter, NullLocator
+
+from tinyllm.results import PAGED_ROW, n, need
+
+OURS, HF, CONTEXT = "#3987e5", "#e0602d", "#8b949e"
+INK = "#6e7681"         # ≥ 4:1 on both white and GitHub's dark #0d1117
+RC = {
+    "font.family": "DejaVu Sans", "font.size": 10, "svg.fonttype": "path", "svg.hashsalt": "tinyllm",
+    "text.color": INK, "axes.labelcolor": INK, "xtick.color": INK, "ytick.color": INK,
+    "axes.edgecolor": CONTEXT, "axes.linewidth": 0.8, "axes.spines.top": False, "axes.spines.right": False,
+    "axes.facecolor": "none", "figure.facecolor": "none", "savefig.transparent": True,
+    "xtick.major.size": 3, "ytick.major.size": 0, "legend.frameon": False,
+}
+GRID = {"color": CONTEXT, "alpha": 0.3, "linewidth": 0.6}
+
+
+def _svg(fig: Figure) -> str:
+    buf = io.StringIO()
+    fig.savefig(buf, format="svg", metadata={"Date": None}, bbox_inches="tight", pad_inches=0.12)
+    return buf.getvalue()
+
+
+def _head(ax, title: str, subtitle: str, lift: float = 0) -> None:
+    """Left-aligned title and subtitle above the plot area, placed in points so they never overlap it."""
+    ax.annotate(title, (0, 1), xycoords="axes fraction", xytext=(0, 30 + lift), textcoords="offset points",
+                fontsize=12.5, fontweight="bold", va="bottom")
+    ax.annotate(subtitle, (0, 1), xycoords="axes fraction", xytext=(0, 12 + lift), textcoords="offset points",
+                fontsize=9.5, va="bottom")
+
+
+def _foot(ax, text: str, below: float = 34) -> None:
+    ax.annotate(text, (0, 0), xycoords="axes fraction", xytext=(0, -below), textcoords="offset points",
+                fontsize=8, va="top")
+
+
+def _source(data: dict, name: str) -> str:
+    commit = (((data[name].get("provenance") or {}).get("git") or {}).get("commit") or "?")[:7]
+    return f"tiny-llm · {data.get('_run_dir', 'runs/bench/v1b')} · commit {commit} · RTX 3070 Ti Laptop GPU"
+
+
+def decode_b1(data: dict) -> Figure:
+    hf = data["tl-hf"]
+    ceiling = hf["bandwidth_bytes_s"] / hf["weight_bytes"]
+    engines = [("tl-kv-graph", "Ours: KV cache + CUDA graphs", OURS), ("tl-kv", "Ours: KV cache", CONTEXT),
+               ("tl-hf", "HF Transformers generate", HF), ("tl-naive", "Ours: naive, no cache", CONTEXT)]
+    rows = [(label, need(data[nm]["rows"], nm, batch=1, prompt_len=128), c) for nm, label, c in engines]
+    rows.sort(key=lambda r: -r[1]["decode_tok_s"])
+    fig = Figure(figsize=(7.6, 2.5))
+    ax = fig.add_subplot()
+    y = range(len(rows))
+    ax.barh(y, [r["decode_tok_s"] for _, r, _ in rows], height=0.62, color=[c for *_, c in rows])
+    ax.set_yticks(list(y), [label for label, *_ in rows])
+    ax.invert_yaxis()
+    for i, (_, r, _) in enumerate(rows):
+        ax.text(r["decode_tok_s"] + ceiling * 0.015, i, f"{n(r['decode_tok_s'])} tok/s · {r['pct_of_ceiling']:.0f}%",
+                va="center", fontsize=9.5)
+    ax.axvline(ceiling, color=INK, linewidth=1.2, linestyle=(0, (4, 3)))
+    ax.text(ceiling * 1.015, len(rows) - 1, f"bandwidth\nceiling\n{n(ceiling)} tok/s", va="center", fontsize=9)
+    ax.set_xlim(0, ceiling * 1.28)
+    ax.xaxis.set_major_locator(FixedLocator([0, 50, 100, 150, 200]))
+    ax.grid(axis="x", **GRID)
+    ax.set_axisbelow(True)
+    ax.spines["left"].set_visible(False)
+    ax.set_xlabel("decode tokens / s at batch 1")
+    _head(ax, "TinyLlama 1.1B decode speed, batch 1",
+          f"128-token prompt, 256 new tokens. % = share of the {n(ceiling)} tok/s ceiling "
+          "(bandwidth ÷ weight bytes).")
+    _foot(ax, _source(data, "tl-kv-graph"), below=40)
+    return fig
+
+
+def batch_scaling(data: dict) -> Figure:
+    hf = data["tl-hf"]
+    ceiling = hf["bandwidth_bytes_s"] / hf["weight_bytes"]
+    batches = sorted({r["batch"] for r in data["tl-kv-graph"]["rows"]} & {r["batch"] for r in hf["rows"]})
+    tok = {nm: [need(data[nm]["rows"], nm, batch=b, prompt_len=128)["decode_tok_s"] for b in batches]
+           for nm in ("tl-kv-graph", "tl-hf")}
+    fig = Figure(figsize=(7.6, 3.6))
+    ax = fig.add_subplot()
+    ax.plot(batches, [ceiling * b for b in batches], color=INK, linewidth=1.2, linestyle=(0, (4, 3)),
+            label=f"bandwidth ceiling ({n(ceiling)} tok/s × batch)")
+    ax.plot(batches, tok["tl-kv-graph"], color=OURS, linewidth=2, marker="o", markersize=7,
+            label="Ours: KV cache + CUDA graphs")
+    ax.plot(batches, tok["tl-hf"], color=HF, linewidth=2, marker="s", markersize=6.5,
+            label="HF Transformers generate")
+    # Every line rises to the right: just left of our markers (incoming line below, ceiling well above) and
+    # below-right of HF's are clear.
+    for nm, off, ha, va in (("tl-kv-graph", (-8, 2), "right", "bottom"), ("tl-hf", (6, -7), "left", "top")):
+        for b, v in zip(batches, tok[nm]):
+            ax.annotate(n(v), (b, v), xytext=off, textcoords="offset points", fontsize=9, ha=ha, va=va)
+    ax.set_xscale("log", base=2)
+    ax.set_yscale("log")
+    ax.xaxis.set_major_locator(FixedLocator(batches))
+    ax.xaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:.0f}"))
+    ax.xaxis.set_minor_locator(NullLocator())
+    ticks = [t for t in (25, 50, 100, 200, 400, 800, 1600, 3200, 6400)
+             if min(tok["tl-hf"][0], tok["tl-kv-graph"][0]) / 2 <= t <= ceiling * batches[-1] * 1.5]
+    ax.yaxis.set_major_locator(FixedLocator(ticks))
+    ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: n(v)))
+    ax.yaxis.set_minor_locator(NullLocator())
+    ax.set_xlim(batches[0] / 1.7, batches[-1] * 1.9)
+    ax.set_ylim(min(tok["tl-hf"][0], tok["tl-kv-graph"][0]) / 1.8, ceiling * batches[-1] * 1.4)
+    ax.grid(axis="y", **GRID)
+    ax.set_xlabel("batch size (sequences decoded together)")
+    ax.set_ylabel("decode tokens / s (log scale)")
+    ax.legend(loc="upper left", fontsize=9)
+    wmib = hf["weight_bytes"] / 2**20
+    _head(ax, "Decode throughput vs batch size, TinyLlama 1.1B",
+          f"Each decode step reads all {n(wmib)} MiB of weights once; a batch shares that read. 128-token prompt.")
+    _foot(ax, _source(data, "tl-kv-graph"), below=40)
+    return fig
+
+
+def train_loss(curve: dict) -> Figure:
+    tps = curve["tokens_per_step"]
+    tr, va = curve["train"], curve["val"]
+    top = 3.2
+    fig = Figure(figsize=(7.6, 3.2))
+    ax = fig.add_subplot()
+    ax.plot([s * tps / 1e6 for s, _ in tr], [v for _, v in tr], color=CONTEXT, linewidth=0.9, label="train loss")
+    ax.plot([s * tps / 1e6 for s, _ in va], [v for _, v in va], color=OURS, linewidth=2, marker="o",
+            markersize=3.5, label="validation loss")
+    last_s, last_v = va[-1]
+    ax.annotate(f"val {last_v:.4f}", (last_s * tps / 1e6, last_v), xytext=(-4, 14), textcoords="offset points",
+                ha="right", fontsize=9.5)
+    ax.annotate(f"← loss starts at {va[0][1]:.1f}, above the chart", (0.06, 0.97), xycoords="axes fraction",
+                fontsize=8.5, va="top")
+    ax.set_ylim(1.0, top)
+    ax.set_xlim(0, tr[-1][0] * tps / 1e6 * 1.01)
+    ax.grid(axis="y", **GRID)
+    ax.set_xlabel("training tokens (millions)")
+    ax.set_ylabel("cross-entropy loss")
+    ax.legend(loc="upper right", fontsize=9)
+    every = va[1][0] - va[0][0] if len(va) > 1 else 0
+    _head(ax, "Training our 30M-parameter model from scratch on TinyStories",
+          f"{curve['params'] / 1e6:.1f}M parameters, {tr[-1][0] * tps / 1e6:,.0f}M tokens"
+          + (f"; validation every {every} steps ({n(every * tps / 1e6)}M tokens)" if every else "") + ".")
+    _foot(ax, f"tiny-llm · docs/data/train_curve.json (from {curve.get('source', 'the training log')})", below=40)
+    return fig
+
+
+def decode_profile(data: dict) -> Figure:
+    d = data["tl-profile"]
+    modes = [("eager", "Eager: one launch per kernel"), ("graph", "CUDA graph: one launch per step")]
+    fig = Figure(figsize=(7.6, 2.1))
+    ax = fig.add_subplot()
+    for i, (mode, _) in enumerate(modes):
+        m = d["modes"][mode]
+        busy = 1 - m["gpu_idle_frac"]
+        ax.barh(i, busy, height=0.58, color=OURS, label="GPU busy" if i == 0 else None)
+        ax.barh(i, m["gpu_idle_frac"], left=busy, height=0.58, color=CONTEXT, alpha=0.55,
+                label="GPU idle, waiting for launches" if i == 0 else None)
+        ops = m.get("gpu_ops_per_step")
+        ax.text(1.02, i, f"{m['ms_per_step']:.1f} ms / step, idle {m['gpu_idle_frac']:.0%}"
+                + (f"\n{n(ops)} kernels" if ops else ""), va="center", fontsize=9)
+    ax.set_yticks(range(len(modes)), [label for _, label in modes])
+    ax.invert_yaxis()
+    ax.set_xlim(0, 1)
+    ax.xaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:.0%}"))
+    ax.spines["left"].set_visible(False)
+    ax.set_xlabel("share of one decode step")
+    ax.legend(loc="lower left", bbox_to_anchor=(0, 1.0), ncols=2, fontsize=9, borderaxespad=0.2)
+    speed = d["modes"]["eager"]["ms_per_step"] / d["modes"]["graph"]["ms_per_step"]
+    _head(ax, "Why CUDA graphs help: TinyLlama 1.1B decode step, batch 1",
+          f"In eager mode Python launches every kernel and the GPU waits; a graph replays the whole step "
+          f"at once ({speed:.2f}× faster).", lift=16)                           # above the legend
+    _foot(ax, _source(data, "tl-profile"), below=40)
+    return fig
+
+
+def serving(data: dict) -> Figure:
+    tls, ms = data["tl-serving"]["rows"], data["30m-serving"]["rows"]
+    st = need(tls, "tl-serving", engine="static", size=8)
+    left = [("Static batching", st), ("Continuous batching", need(tls, "tl-serving", engine="continuous", size=8)),
+            ("Continuous + CUDA graphs", need(tls, "tl-serving", engine="cont+graph", size=8))]
+    c8 = need(ms, "30m-serving", engine="continuous", size=8)
+    pg = need(ms, "30m-serving", engine=PAGED_ROW[0], size=PAGED_ROW[1])
+    right = [(f"Contiguous KV: {c8['peak_active']} sequences", c8), (f"Paged KV: {pg['peak_active']} sequences", pg)]
+    fig = Figure(figsize=(7.6, 2.4))
+    axes = fig.subplots(1, 2, gridspec_kw={"wspace": 1.05, "width_ratios": [1, 1]})
+    for ax, rows, title in ((axes[0], left, "TinyLlama 1.1B, 8 slots"),
+                            (axes[1], right, f"30M model, same {c8['kv_reserved_mib']:.0f} MiB of KV")):
+        base = rows[0][1]["tok_s"]
+        vals = [r["tok_s"] for _, r in rows]
+        ax.barh(range(len(rows)), vals, height=0.6, color=OURS)
+        ax.set_yticks(range(len(rows)), [label for label, _ in rows])
+        ax.invert_yaxis()
+        for i, v in enumerate(vals):
+            ax.text(v + max(vals) * 0.03, i, n(v) + (f" ({v / base:.2f}×)" if i else ""), va="center", fontsize=9)
+        ax.set_xlim(0, max(vals) * 1.45)
+        ax.xaxis.set_major_locator(FixedLocator([]))
+        ax.spines["left"].set_visible(False)
+        ax.spines["bottom"].set_visible(False)
+        ax.set_title(title, loc="left", fontsize=10, color=INK)
+    _head(axes[0], f"Serving throughput, {data['tl-serving']['workload']['requests']} mixed-length requests",
+          "Tokens / s; × = vs the first row. Continuous batching refills a slot as soon as its request ends; "
+          "the paged cache fits more sequences in the same memory.", lift=18)   # above the panel titles
+    _foot(axes[0], _source(data, "tl-serving"), below=12)
+    return fig
+
+
+AQUA = "#1aa775"        # third engine; blue / orange / aqua checked for colour-vision deficiency on both themes
+LOAD_STYLE = {"Ours": (OURS, "o", 7), "vLLM": (HF, "s", 6.5), "llama.cpp": (AQUA, "D", 6)}
+
+
+def load_sweep(load: dict) -> Figure:
+    """docs/data/load_sweep.json (scripts/extract_load_sweep.py): the 4.8a equal-seat sweep, one row per rate."""
+    col = {c: i for i, c in enumerate(load["columns"])}
+    eng = {name: [[p[col[c]] for p in e["points"]] for c in ("rate_rps", "achieved_rps", "ttft_p95_s", "shed_rate")]
+           for name, e in load["engines"].items()}
+    top = max(r[-1] for r, *_ in eng.values())
+    fig = Figure(figsize=(7.6, 3.3))
+    left, right = fig.subplots(1, 2, gridspec_kw={"wspace": 0.7})
+    peak = max(max(a) for _, a, *_ in eng.values())
+    left.plot([0, peak * 1.12], [0, peak * 1.12], color=INK, linewidth=1.2, linestyle=(0, (4, 3)))
+    left.annotate("every request\nserved", (peak, peak), xytext=(-8, 0), textcoords="offset points", ha="right",
+                  va="center", fontsize=8.5)
+    slo = load["slo_ttft_p95_s"]
+    right.axhline(slo, color=INK, linewidth=1.2, linestyle=(0, (4, 3)))
+    right.annotate(f"SLO {n(slo)} s", (top, slo), xytext=(0, 3), textcoords="offset points", ha="right", va="bottom",
+                   fontsize=8.5)
+    for name, (rate, done, ttft, shed) in eng.items():
+        c, m, ms = LOAD_STYLE[name]
+        for ax, y in ((left, done), (right, ttft)):
+            ax.plot(rate, y, color=c, linewidth=2, marker=m, markersize=ms)
+        left.annotate(f"{name} {done[-1]:.1f}", (rate[-1], done[-1]), xytext=(7, 0), textcoords="offset points",
+                      va="center", fontsize=9, annotation_clip=False)
+        if shed[-1] > 0:        # the gap between our flat line and the diagonal is rejected, not queued
+            i = len(rate) - 2
+            left.annotate(f"{shed[i] * 100:.0f}–{shed[-1]:.0%} rejected\nwith a 429", (rate[-1], done[-1]), xytext=(0, 9),
+                          textcoords="offset points", ha="right", va="bottom", fontsize=8.5)
+        end = f"{name} {ttft[-1]:.0f} s" if ttft[-1] >= 10 else f"{name} {ttft[-1]:.1f} s"
+        if shed[-1] > 0:        # flat because the bounded queue rejects, not because the engine is faster
+            right.annotate(end + ": capped;\na full queue answers\nnew requests with a 429", (rate[-1], ttft[-1]),
+                           xytext=(7, 6), textcoords="offset points", va="top", fontsize=9, annotation_clip=False)
+        else:
+            right.annotate(end, (rate[-1], ttft[-1]), xytext=(7, 0), textcoords="offset points", va="center",
+                           fontsize=9, annotation_clip=False)
+    for ax in (left, right):
+        ax.set_xlim(0, top * 1.04)
+        ax.xaxis.set_major_locator(FixedLocator([0, 4, 8, 12, 16]))
+        ax.grid(axis="y", **GRID)
+        ax.set_xlabel("offered requests / s")
+    left.set_ylim(0, peak * 1.12)
+    left.set_title("Completed requests / s", loc="left", fontsize=10, color=INK)
+    right.set_yscale("log")
+    right.yaxis.set_major_locator(FixedLocator([0.1, 1, 10, 100]))
+    right.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g} s"))
+    right.yaxis.set_minor_locator(NullLocator())
+    right.set_ylim(0.02, max(max(t) for _, _, t, _ in eng.values()) * 2)
+    right.set_title("Time to first token p95, served requests", loc="left", fontsize=10, color=INK)
+    w = load["workload"]
+    _head(left, "TinyLlama 1.1B under load: 8 sequences per engine",
+          f"Poisson arrivals, {w['max_tokens']} new tokens per request, {w['duration']:.0f} s per rate. Past capacity "
+          "we reject the excess; vLLM and llama.cpp queue it.", lift=18)       # above the panel titles
+    commit = next(iter(load["engines"].values()))["commit"]
+    _foot(left, f"tiny-llm · docs/data/load_sweep.json ({load['source']}) · commit {commit} · "
+                "RTX 3070 Ti Laptop GPU, thermally throttled", below=40)
+    return fig
+
+
+DRAW = {"decode_b1.svg": decode_b1, "batch_scaling.svg": batch_scaling,
+        "decode_profile.svg": decode_profile, "serving.svg": serving}
+
+
+def render_plots(data: dict, curve: dict | None, load: dict | None = None) -> dict[str, str]:
+    """{file name: SVG text}; train_loss.svg and load_sweep.svg only when their docs/data files exist."""
+    with matplotlib.rc_context(RC):
+        out = {name: _svg(draw(data)) for name, draw in DRAW.items()}
+        if curve is not None:
+            out["train_loss.svg"] = _svg(train_loss(curve))
+        if load is not None:
+            out["load_sweep.svg"] = _svg(load_sweep(load))
+    return out

src/tinyllm/provenance.py

@@ -0,0 +1,237 @@
+"""Step 4.12 (D-025): record where every benchmark number came from.
+
+    prov = Provenance(weights="fp32 + bf16 autocast").start()   # after model load, before timing
+    ... timed runs ...
+    json_out["provenance"] = prov.finish()
+
+Recorded:
+  static   UTC time, git commit and whether tracked files are modified, python / torch / CUDA /
+           cuDNN versions, GPU name and driver, plus any fields passed in (e.g. the weight dtype)
+  before   GPU sample, WSL load average, Windows CPU load %
+  during   GPU samples from one `nvidia-smi -lms` process running next to the benchmark:
+           min / median / max of SM clock, memory clock, temperature, power, utilisation; every
+           throttle reason seen; and the WSL load average
+  after    same as before
+
+Notes:
+  - WSL's /proc/loadavg only sees the Linux VM. Load from Windows programs (the usual cause of
+    the 4.5 outliers) shows up only in the Windows CPU %, read via powershell.exe.
+  - That read costs ~2 s of CPU, so it runs only before and after, never during the timed region.
+  - The GPU sampler is a separate process, so it costs the benchmark's Python process nothing.
+    Every probe is optional: missing tools give None, never an error.
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+import shutil
+import statistics
+import subprocess
+import threading
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+
+import torch
+
+GPU_FIELDS = ["pstate", "clocks.sm", "clocks.mem", "temperature.gpu", "power.draw",
+              "utilization.gpu", "memory.used", "clocks_throttle_reasons.active"]
+NUMERIC = ["clocks.sm", "clocks.mem", "temperature.gpu", "power.draw", "utilization.gpu", "memory.used"]
+BUSY_UTIL = 10  # % GPU utilisation at or above which a sample counts as "busy"
+
+# NVML clocksThrottleReasons bits (nvml.h)
+THROTTLE_BITS = {
+    0x1: "gpu_idle", 0x2: "applications_clocks_setting", 0x4: "sw_power_cap", 0x8: "hw_slowdown",
+    0x10: "sync_boost", 0x20: "sw_thermal_slowdown", 0x40: "hw_thermal_slowdown",
+    0x80: "hw_power_brake_slowdown", 0x100: "display_clock_setting",
+}
+REPO = Path(__file__).resolve().parents[2]
+
+
+def _run(cmd: list[str], timeout: float = 10.0) -> str | None:
+    try:
+        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, cwd=REPO)
+    except (OSError, subprocess.SubprocessError):
+        return None
+    return r.stdout.strip() if r.returncode == 0 else None
+
+
+def git_info() -> dict:
+    commit = _run(["git", "rev-parse", "HEAD"])
+    status = _run(["git", "status", "--porcelain", "--untracked-files=no"])
+    return {"commit": commit, "dirty": None if status is None else bool(status)}
+
+
+def claim_output(path: str | Path, overwrite: bool = False) -> Path:
+    """Call before a benchmark starts: refuse an existing result file unless overwrite (6.8).
+
+    The 11-04 TinyLlama serving JSON was lost because a later run with other flags wrote the same
+    default name. Checking at startup costs nothing; finding out after an hour's run costs the hour.
+    """
+    path = Path(path)
+    if path.exists() and not overwrite:
+        raise SystemExit(f"{path} already exists: pass a new --out, or --overwrite to replace it")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def decode_throttle(mask: int | None) -> list[str]:
+    if mask is None:
+        return []
+    names = [name for bit, name in THROTTLE_BITS.items() if mask & bit]
+    rest = mask & ~sum(THROTTLE_BITS)
+    return names + ([f"unknown_0x{rest:x}"] if rest else [])
+
+
+def parse_gpu_line(line: str) -> dict | None:
+    """One `nvidia-smi --query-gpu=<GPU_FIELDS> --format=csv,noheader,nounits` line -> dict."""
+    parts = [p.strip() for p in line.strip().split(",")]
+    if len(parts) != len(GPU_FIELDS):
+        return None
+    out: dict = {}
+    for k, v in zip(GPU_FIELDS, parts):
+        if v in ("[N/A]", "N/A", "[Not Supported]", ""):
+            out[k] = None
+        elif k == "pstate":
+            out[k] = v
+        elif k == "clocks_throttle_reasons.active":
+            out[k] = int(v, 16)
+        else:
+            try:
+                out[k] = float(v)
+            except ValueError:
+                out[k] = None
+    return out
+
+
+def _smi() -> str | None:
+    return shutil.which("nvidia-smi")
+
+
+def gpu_sample() -> dict | None:
+    smi = _smi()
+    if smi is None:
+        return None
+    out = _run([smi, f"--query-gpu={','.join(GPU_FIELDS)}", "--format=csv,noheader,nounits"])
+    return parse_gpu_line(out.splitlines()[0]) if out else None
+
+
+def wsl_loadavg() -> float | None:
+    try:
+        return float(Path("/proc/loadavg").read_text().split()[0])
+    except (OSError, ValueError, IndexError):
+        return None
+
+
+def windows_cpu_percent() -> float | None:
+    ps = shutil.which("powershell.exe")
+    if ps is None:
+        return None
+    out = _run([ps, "-NoProfile", "-Command",
+                "(Get-CimInstance Win32_Processor | Measure-Object -Property LoadPercentage -Average).Average"],
+               timeout=20)
+    try:
+        return float(out) if out else None
+    except ValueError:
+        return None
+
+
+def host_sample() -> dict:
+    return {"gpu": gpu_sample(), "wsl_loadavg_1m": wsl_loadavg(), "windows_cpu_percent": windows_cpu_percent()}
+
+
+def summarize(samples: list[dict], loads: list[float]) -> dict:
+    """min / median / max per numeric GPU field, the throttle reasons seen, and the WSL load."""
+    out: dict = {"n_samples": len(samples)}
+    for k in NUMERIC:
+        vals = [s[k] for s in samples if s.get(k) is not None]
+        out[k] = ({"min": min(vals), "median": statistics.median(vals), "max": max(vals)} if vals else None)
+    # Idle gaps (startup, CPU-only work between runs) drag the minimum down; this is the clock while working.
+    busy = [s["clocks.sm"] for s in samples
+            if s.get("clocks.sm") is not None and (s.get("utilization.gpu") or 0) >= BUSY_UTIL]
+    out["clocks.sm_busy"] = ({"min": min(busy), "median": statistics.median(busy), "max": max(busy),
+                              "n": len(busy)} if busy else None)
+    mask = 0
+    for s in samples:
+        mask |= s.get("clocks_throttle_reasons.active") or 0
+    out["throttle_reasons_seen"] = decode_throttle(mask)
+    out["pstates_seen"] = sorted({s["pstate"] for s in samples if s.get("pstate")})
+    out["wsl_loadavg_1m"] = ({"min": min(loads), "max": max(loads)} if loads else None)
+    return out
+
+
+class GpuMonitor:
+    """Streams `nvidia-smi -lms <interval>` in a background thread (one process for the whole run)."""
+
+    def __init__(self, interval_ms: int = 500):
+        self.interval_ms = interval_ms
+        self.samples: list[dict] = []
+        self.loads: list[float] = []
+        self._proc: subprocess.Popen | None = None
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> GpuMonitor:
+        smi = _smi()
+        if smi is None:
+            return self
+        try:
+            self._proc = subprocess.Popen(
+                [smi, f"--query-gpu={','.join(GPU_FIELDS)}", "--format=csv,noheader,nounits",
+                 f"-lms={self.interval_ms}"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True)
+        except OSError:
+            return self
+        self._thread = threading.Thread(target=self._read, daemon=True)
+        self._thread.start()
+        return self
+
+    def _read(self) -> None:
+        assert self._proc is not None and self._proc.stdout is not None
+        for line in self._proc.stdout:
+            s = parse_gpu_line(line)
+            if s is not None:
+                self.samples.append(s)
+                load = wsl_loadavg()
+                if load is not None:
+                    self.loads.append(load)
+
+    def stop(self) -> dict:
+        if self._proc is not None:
+            self._proc.terminate()
+            try:
+                self._proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self._proc.kill()
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+        return summarize(self.samples, self.loads)
+
+
+class Provenance:
+    def __init__(self, interval_ms: int = 500, **extra):
+        self.record: dict = {
+            "time_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+            "git": git_info(),
+            "python": platform.python_version(),
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "cudnn": torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else None,
+            "gpu_name": torch.cuda.get_device_name() if torch.cuda.is_available() else None,
+            "driver": _run([_smi(), "--query-gpu=driver_version", "--format=csv,noheader"]) if _smi() else None,
+            "cpu_count": os.cpu_count(),
+            **extra,
+        }
+        self._monitor = GpuMonitor(interval_ms)
+        self._t0 = 0.0
+
+    def start(self) -> Provenance:
+        self.record["before"] = host_sample()
+        self._monitor.start()
+        self._t0 = time.perf_counter()
+        return self
+
+    def finish(self) -> dict:
+        self.record["monitored_s"] = round(time.perf_counter() - self._t0, 2)
+        self.record["during"] = self._monitor.stop()
+        self.record["after"] = host_sample()
+        return self.record

src/tinyllm/quant.py

@@ -0,0 +1,139 @@
+"""Step 4.6: INT8 weight-only quantization (per-output-channel, symmetric absmax).
+
+Each Linear weight row j gets one scale s_j = max|W_j| / 127 and is stored as
+q_j = round(W_j / s_j), clipped to [-127, 127] (int8, 1 byte instead of 2 for bf16).
+Activations stay in floating point ("weight-only", W8A16).
+
+Because the scale is per output channel, it factors out of the matmul exactly:
+    y_j = sum_i x_i * (q_ji * s_j) = s_j * sum_i x_i * q_ji
+so the forward is  F.linear(x, q.to(dtype)) * s : cast int8 -> bf16 (exact: |q| <= 127 < 256),
+matmul, scale the output. Values are the same as dequantizing W first. In PyTorch the int8 ->
+bf16 cast is its own kernel (backend "torch"). Backend "triton" (4.10b) runs one kernel that reads
+the int8 weights directly (tinyllm.kernels.int8_gemv) for decode-sized inputs on CUDA, and falls
+back to the torch path otherwise (CPU, fp32, prefill-sized inputs, no Triton).
+
+Embeddings stay in floating point: a lookup reads only the B rows it needs, not the table.
+With tied embeddings, lm_head gets its own int8 copy and the embedding keeps the original.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+QMAX = 127  # symmetric: -128 is unused so that +/- ranges match
+BACKENDS = ("torch", "triton")
+
+
+def quantize_per_channel(w: Tensor) -> tuple[Tensor, Tensor]:
+    """(out, in) float weight -> (int8 q (out, in), fp32 scale (out,)). An all-zero row gets scale 0."""
+    w = w.detach().float()
+    scale = w.abs().amax(dim=1) / QMAX
+    q = torch.round(w / scale.clamp_min(1e-12)[:, None]).clamp_(-QMAX, QMAX).to(torch.int8)
+    return q, scale
+
+
+def dequantize(q: Tensor, scale: Tensor) -> Tensor:
+    return q.float() * scale[:, None]
+
+
+class Int8Linear(nn.Module):
+    """Drop-in for a bias-free nn.Linear, with int8 weights and one fp32 scale per output channel."""
+
+    def __init__(self, q: Tensor, scale: Tensor, backend: str = "torch"):
+        super().__init__()
+        self.in_features, self.out_features = q.shape[1], q.shape[0]
+        self.register_buffer("qweight", q)
+        self.register_buffer("scale", scale)
+        self.backend = _check_backend(backend)
+
+    @classmethod
+    def from_linear(cls, m: nn.Linear, backend: str = "torch") -> Int8Linear:
+        if m.bias is not None:
+            raise ValueError("Int8Linear supports bias-free Linear only (the model has no biases)")
+        return cls(*quantize_per_channel(m.weight), backend=backend)
+
+    def forward(self, x: Tensor) -> Tensor:
+        dev = x.device.type
+        dt = torch.get_autocast_dtype(dev) if torch.is_autocast_enabled(dev) else x.dtype
+        if self.backend == "triton" and x.is_cuda and dt in (torch.bfloat16, torch.float16):
+            kernel = _triton_kernel()
+            rows = x.numel() // self.in_features
+            if kernel is not None and rows <= kernel.MAX_M:
+                y = kernel.int8_matmul(x.reshape(rows, self.in_features).to(dt), self.qweight, self.scale)
+                return y.view(*x.shape[:-1], self.out_features)
+        # Output keeps fp32 under autocast (bf16 matmul result * fp32 scale), so the scale isn't
+        # rounded to bf16. Downstream ops cast as they would after an nn.Linear.
+        return F.linear(x.to(dt), self.qweight.to(dt)) * self.scale
+
+    def extra_repr(self) -> str:
+        return (f"in_features={self.in_features}, out_features={self.out_features}, int8 per-channel, "
+                f"backend={self.backend}")
+
+
+def _check_backend(backend: str) -> str:
+    if backend not in BACKENDS:
+        raise ValueError(f"unknown Int8Linear backend {backend!r}; one of {BACKENDS}")
+    return backend
+
+
+_KERNEL: list = []
+
+
+def _triton_kernel():
+    """tinyllm.kernels.int8_gemv, or None without Triton (imported once, on first CUDA use)."""
+    if not _KERNEL:
+        try:
+            from tinyllm.kernels import int8_gemv
+        except ImportError:
+            int8_gemv = None
+        _KERNEL.append(int8_gemv)
+    return _KERNEL[0]
+
+
+def quantize_model(model: nn.Module, backend: str = "torch") -> int:
+    """Replace every nn.Linear in `model` (in place) with an Int8Linear. Returns how many."""
+    targets = [(parent, name, child) for parent in model.modules()
+               for name, child in parent.named_children() if isinstance(child, nn.Linear)]
+    for parent, name, child in targets:
+        q = Int8Linear.from_linear(child, backend).to(child.weight.device)
+        setattr(parent, name, q)
+    return len(targets)
+
+
+def set_backend(model: nn.Module, backend: str) -> int:
+    """Switch every Int8Linear in `model` to `backend` (same weights, no copy). Returns how many."""
+    _check_backend(backend)
+    mods = [m for m in model.modules() if isinstance(m, Int8Linear)]
+    for m in mods:
+        m.backend = backend
+    return len(mods)
+
+
+def weight_bytes(model: nn.Module, float_bytes: int | None = None) -> dict[str, int]:
+    """Bytes of persistent weights (params + saved buffers; RoPE tables excluded), split into the
+    matmul weights read every decode step and the embedding table (only B rows read per step).
+
+    float_bytes: count float tensors at this width (2 = as if stored in bf16); None = as stored.
+    Int8Linear tensors are always counted as stored (int8 weights, fp32 scales).
+    A tied matrix is counted once, under matmul: lm_head reads all of it every step."""
+    def nbytes(t: Tensor, as_stored: bool) -> int:
+        if t.is_floating_point() and float_bytes is not None and not as_stored:
+            return t.numel() * float_bytes
+        return t.numel() * t.element_size()
+
+    is_mm = lambda m: isinstance(m, (nn.Linear, Int8Linear))  # noqa: E731
+    mods = sorted(model.modules(), key=lambda m: not is_mm(m))  # matmul modules claim shared tensors first
+    seen, out = set(), {"matmul": 0, "embedding": 0, "other": 0}
+    for mod in mods:
+        tensors = list(mod.named_parameters(recurse=False)) + [
+            (n, b) for n, b in mod.named_buffers(recurse=False)
+            if n not in mod._non_persistent_buffers_set]
+        for _, t in tensors:
+            if id(t) in seen:
+                continue
+            seen.add(id(t))
+            kind = "matmul" if is_mm(mod) else "embedding" if isinstance(mod, nn.Embedding) else "other"
+            out[kind] += nbytes(t, as_stored=isinstance(mod, Int8Linear))  # its scales stay fp32
+    return out

src/tinyllm/results.py

@@ -0,0 +1,654 @@
+"""Step 6.8 (D-029): the results tables, rendered from one benchmark run's JSON files.
+
+scripts/bench_all.py runs every benchmark in BENCHES into one directory (runs/bench/v1b/), then calls
+render_results() for docs/RESULTS.md and render_readme() for the README's <!-- bench:NAME --> blocks.
+Everything here is pure (JSON in, markdown out), so the CPU tests cover it without a GPU.
+
+Rules this module enforces:
+  - every table is built from one run directory: nothing is mixed in from other sessions
+  - a missing file is an error naming the `--only` command that makes it, never an empty table
+  - each table carries a heat line from its files' provenance, since absolute tok/s on this laptop
+    moves 10-20% between sessions with temperature
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+# A sustained clock below this, or a hardware slowdown, means the numbers are not this GPU's normal
+# (11-04 runs: busy median ~1,600 MHz at 85 °C, with sw_power_cap / sw_thermal_slowdown only).
+MIN_CLOCK_MHZ = 1400
+HARD_THROTTLE = {"hw_slowdown", "hw_thermal_slowdown", "hw_power_brake_slowdown"}
+PAGED_ROW = ("paged/16", 32)        # the README's paged row: 16-token blocks, 32 slots, the 8-slot budget
+
+
+@dataclass(frozen=True)
+class Bench:
+    name: str
+    script: str
+    args: tuple[str, ...]
+    out: str                         # path inside the run directory
+    minutes: int                     # rough, for the run plan printout
+    env: tuple[tuple[str, str], ...] = ()   # extra environment for this benchmark, recorded in the manifest
+
+    @property
+    def out_dir_flag(self) -> bool:  # profile_decode takes a directory, the others a file
+        return self.script == "profile_decode.py"
+
+
+TL = ("--model", "tinyllama")
+# The naive engine runs a full forward per token, so its activations and logits grow by one position every
+# step. No cached block fits the next request, the caching allocator keeps reserving new ones until the GPU
+# is full (7,991 MiB in use for a 475 MiB peak at 30M B=32 P=128), and a cuBLAS call then fails
+# (6.8, 2026-11-06). Expandable segments let a block grow in place: 695 MiB, same speed. Only the naive
+# benchmarks need it; the other engines use fixed-size caches.
+GROWING = (("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True"),)
+BENCHES = [
+    Bench("30m-naive", "bench_generate.py", ("--engine", "naive", "--dtype", "bf16"), "30m/naive.json", 4, GROWING),
+    Bench("30m-kv", "bench_generate.py", ("--engine", "kv", "--dtype", "bf16"), "30m/kv.json", 2),
+    Bench("30m-kv-graph", "bench_generate.py", ("--engine", "kv-graph", "--dtype", "bf16"), "30m/kv-graph.json", 2),
+    Bench("30m-serving", "bench_serving.py", ("--cuda-graph", "--paged"), "30m/serving.json", 12),
+    Bench("30m-int8", "bench_quant.py", (), "30m/int8.json", 4),
+    Bench("30m-profile", "profile_decode.py", ("--batch", "1"), "30m/profile/decode_main_B1.json", 1),
+    Bench("tl-logits", "tinyllama_vs_hf.py", (), "tinyllama/logits_vs_hf.json", 4),
+    Bench("tl-hf", "bench_generate.py", (*TL, "--engine", "hf"), "tinyllama/hf.json", 12),
+    Bench("tl-naive", "bench_generate.py", (*TL, "--engine", "naive"), "tinyllama/naive.json", 6, GROWING),
+    Bench("tl-kv", "bench_generate.py", (*TL, "--engine", "kv"), "tinyllama/kv.json", 10),
+    Bench("tl-kv-graph", "bench_generate.py", (*TL, "--engine", "kv-graph"), "tinyllama/kv-graph.json", 6),
+    Bench("tl-serving", "bench_serving.py", (*TL, "--cuda-graph"), "tinyllama/serving.json", 15),
+    Bench("tl-int8", "bench_quant.py", TL, "tinyllama/int8.json", 10),
+    # 4.10b: the Triton INT8 kernel, end to end with CUDA graphs, and its matmuls alone
+    Bench("tl-int8-graph", "bench_quant.py", (*TL, "--engine", "kv-graph", "--triton"),
+          "tinyllama/int8_graph_triton.json", 12),
+    Bench("tl-int8-gemv", "bench_int8_gemv.py", (), "tinyllama/int8_gemv.json", 3),
+    # 4.10a: the Triton RMSNorm + residual add, alone and end to end with CUDA graphs
+    Bench("tl-rmsnorm", "bench_rmsnorm.py", (), "tinyllama/rmsnorm_graph.json", 10),
+    Bench("tl-profile", "profile_decode.py", (*TL, "--batch", "1"), "tinyllama/profile/decode_tinyllama_B1.json", 2),
+]
+BY_NAME = {b.name: b for b in BENCHES}
+MANIFEST = "manifest.json"
+# Optional results run by hand, read when present (None otherwise). 4.10c: tl-logits with our RMSNorm forced
+# to the PyTorch code (tinyllama_vs_hf.py --norm-backend torch), which shows where D-032's logit gap comes from.
+EXTRAS = {"tl-logits-torchnorm": "tinyllama/logits_vs_hf_torchnorm.json"}
+
+
+class MissingResult(Exception):
+    pass
+
+
+def missing(run_dir: str | Path) -> list[Bench]:
+    return [b for b in BENCHES if not (Path(run_dir) / b.out).exists()]
+
+
+def load_run(run_dir: str | Path) -> dict:
+    """{bench name: its JSON, "manifest": manifest or {}}; raises MissingResult listing every gap."""
+    run_dir = Path(run_dir)
+    gaps = missing(run_dir)
+    if gaps:
+        lines = "\n".join(f"  {run_dir / b.out}  ->  scripts/bench_all.py --only {b.name}" for b in gaps)
+        raise MissingResult(f"{len(gaps)} result file(s) missing from {run_dir}:\n{lines}")
+    data = {b.name: json.loads((run_dir / b.out).read_text()) for b in BENCHES}
+    for name, out in EXTRAS.items():
+        data[name] = json.loads((run_dir / out).read_text()) if (run_dir / out).exists() else None
+    m = run_dir / MANIFEST
+    data["manifest"] = json.loads(m.read_text()) if m.exists() else {}
+    return data
+
+
+# ---------------------------------------------------------------- formatting
+
+
+def n(x: float, d: int = 0) -> str:
+    return f"{x:,.{d}f}"
+
+
+def table(header: list[str], rows: list[list[str]]) -> str:
+    out = ["| " + " | ".join(header) + " |", "|" + "---|" * len(header)]
+    out += ["| " + " | ".join(r) + " |" for r in rows]
+    return "\n".join(out)
+
+
+def row(rows: list[dict], **kv) -> dict | None:
+    return next((r for r in rows if all(r.get(k) == v for k, v in kv.items())), None)
+
+
+def need(rows: list[dict], what: str, **kv) -> dict:
+    r = row(rows, **kv)
+    if r is None or "skipped" in r:
+        raise MissingResult(f"{what}: no row with {kv}")
+    return r
+
+
+# ---------------------------------------------------------------- heat
+
+
+def heat(prov: dict | None) -> dict | None:
+    """Busy SM clock median / min, peak temperature, throttle reasons (gpu_idle dropped), warnings."""
+    if not prov or not prov.get("during"):
+        return None
+    d = prov["during"]
+    clock = d.get("clocks.sm_busy") or d.get("clocks.sm") or {}
+    reasons = [r for r in d.get("throttle_reasons_seen") or [] if r != "gpu_idle"]
+    warn = sorted(HARD_THROTTLE.intersection(reasons))
+    # A low clock alone is not throttling: small or host-bound steps leave the GPU idle and it clocks down by
+    # itself (30M KV decode: 1,095 MHz at 65 °C). Flag it only when the driver also reports a thermal slowdown.
+    if (clock.get("median") is not None and clock["median"] < MIN_CLOCK_MHZ
+            and "sw_thermal_slowdown" in reasons):
+        warn.append(f"busy clock median {clock['median']:.0f} MHz < {MIN_CLOCK_MHZ}")
+    return {"clock_median": clock.get("median"), "clock_min": clock.get("min"),
+            "temp_max": (d.get("temperature.gpu") or {}).get("max"), "reasons": reasons, "warn": warn}
+
+
+def heat_text(h: dict | None) -> str:
+    if h is None:
+        return "not recorded"
+    parts = []
+    if h["clock_median"] is not None:
+        parts.append(f"{n(h['clock_median'])} MHz (min {n(h['clock_min'])})")
+    if h["temp_max"] is not None:
+        parts.append(f"{h['temp_max']:.0f} °C")
+    parts.append(" + ".join(h["reasons"]) or "no throttling")
+    return ("⚠ " if h["warn"] else "") + ", ".join(parts)
+
+
+def heat_line(data: dict, names: list[str], labels: list[str] | None = None) -> str:
+    labels = labels or names
+    body = "; ".join(f"{lab}: {heat_text(heat(data[nm].get('provenance')))}" for nm, lab in zip(names, labels))
+    return f"_Heat (busy SM clock median, peak GPU temperature, throttle reasons): {body}._"
+
+
+def warnings(data: dict) -> list[str]:
+    out = []
+    for b in BENCHES:
+        h = heat(data[b.name].get("provenance"))
+        if h and h["warn"]:
+            out.append(f"{b.name}: {', '.join(h['warn'])} -> consider scripts/bench_all.py --only {b.name}")
+    return out
+
+
+# ---------------------------------------------------------------- tables
+
+
+def ceiling_note(d: dict) -> str:
+    bw, wb = d["bandwidth_bytes_s"], d["weight_bytes"]
+    return (f"{bw / 1e9:.0f} GB/s ÷ {n(wb / 2**20)} MiB = **{n(bw / wb)} tok/s per sequence**")
+
+
+def repeats_note(d: dict) -> str:
+    r = (d.get("provenance") or {}).get("repeats")
+    return f"median of {r} runs" if r else "median of repeated runs"
+
+
+TL_ENGINES = [("tl-hf", "HF Transformers `generate` (bf16, SDPA)"),
+              ("tl-naive", "Ours, naive (full forward per token)"),
+              ("tl-kv", "Ours, KV cache"),
+              ("tl-kv-graph", "Ours, KV cache + CUDA graphs")]
+
+
+def tl_b1_table(data: dict, prompt_len: int = 128) -> str:
+    rows = []
+    for name, label in TL_ENGINES:
+        r = need(data[name]["rows"], name, batch=1, prompt_len=prompt_len)
+        cells = [label, n(r["decode_tok_s"]), f"{r['pct_of_ceiling']:.0f}%", f"{r['ms_per_step']:.1f}"]
+        rows.append([f"**{c}**" for c in cells] if name == "tl-kv-graph" else cells)
+    return table(["Engine", "tok/s", "% of ceiling", "ms / step"], rows)
+
+
+def decode_grid(data: dict, engines: list[tuple[str, str]]) -> str:
+    plens = sorted({r["prompt_len"] for nm, _ in engines for r in data[nm]["rows"]})
+    out, slides = [], False
+    for name, label in engines:
+        rows = data[name]["rows"]
+        for b in sorted({r["batch"] for r in rows}):
+            cells = [label, str(b)]
+            for p in plens:
+                r = row(rows, batch=b, prompt_len=p)
+                if r is None or "skipped" in r:
+                    cells.append("—")
+                    continue
+                mark = " ‡" if r.get("window_slides") else ""
+                slides |= bool(mark)
+                cells.append(f"{n(r['decode_tok_s'])} ({r['pct_of_ceiling']:.0f}%){mark}")
+            out.append(cells)
+    text = table(["Engine", "Batch", *[f"P={p}" for p in plens]], out)
+    text += "\n\nDecode tok/s, with the % of the bandwidth ceiling at that batch size. — = skipped: " \
+            "prompt + new tokens exceed the context, which a KV cache can't slide."
+    if slides:
+        text += " ‡ = prompt + new tokens exceed the context, so naive slides its window."
+    return text
+
+
+def int8_section(d: dict) -> str:
+    base = d["baseline"]
+    variants = list(dict.fromkeys(r["variant"] for r in d["rows"]))
+    others = [v for v in variants if v != base]
+    rows = []
+    for b in sorted({r["batch"] for r in d["rows"]}):
+        tok = {v: need(d["rows"], "int8", variant=v, batch=b)["decode_tok_s"] for v in variants}
+        rows.append([str(b), *[n(tok[v]) for v in variants], *[f"{tok[v] / tok[base]:.2f}×" for v in others]])
+    text = table(["Batch", *[f"{v} tok/s" for v in variants], *[f"{v} / {base}" for v in others]], rows)
+    mib = {v: sum(d["weights_mib"][v].values()) for v in variants}
+    loss, tf, same = d["val_loss"], d["teacher_forced"], d["greedy_rows_identical"]
+    n_seq = same[base]
+    text += "\n\n" + "\n".join([
+        "- Weights: " + ", ".join(f"{v} {n(mib[v])} MiB" for v in variants) + ".",
+        "- Val loss: " + ", ".join(f"{v} {loss[v]:.4f}" for v in variants)
+        + " (" + ", ".join(f"{v} {loss[v] - loss[base]:+.4f}" for v in others) + f" vs {base}).",
+        f"- Teacher-forced on {base}'s greedy tokens: argmax agreement "
+        + ", ".join(f"{v} {tf[v]['argmax_agree']:.2%}" for v in variants)
+        + f" (median top-1 margin {tf['_median_gap']:.2f}).",
+        f"- Greedy rows identical to {base}: " + ", ".join(f"{v} {same[v]}/{n_seq}" for v in variants) + ".",
+    ])
+    return text
+
+
+GEMV_VARIANTS = [("bf16", "bf16"), ("torch", "PyTorch int8"), ("triton", "Triton int8")]
+
+
+def gemv_bandwidth(d: dict, m: int, v: str) -> tuple[float, float]:
+    """(bytes one decode step's matmuls must read, % of the measured bandwidth) for variant v at m rows."""
+    rows = [r for r in d["rows"] if r["m"] == m]
+    nbytes = sum(r[f"{v}_gb_s"] * r[f"{v}_us"] * 1e3 * r["per_step"] for r in rows)    # GB/s × µs = kB
+    return nbytes, 100 * nbytes / (d["step_us"][str(m)][v] * 1e-6) / d["bandwidth"]
+
+
+def gemv_table(d: dict) -> str:
+    ms = sorted(map(int, d["step_us"]))
+    rows = []
+    for m in ms:
+        s = d["step_us"][str(m)]
+        rows.append([str(m), *[f"{s[v] / 1e3:.2f}" for v, _ in GEMV_VARIANTS],
+                     _ratio(s["bf16"] / s["triton"]), _ratio(s["torch"] / s["triton"], 1)])
+    text = table(["Rows (batch)", *[f"{lab} ms" for _, lab in GEMV_VARIANTS], "Triton vs bf16",
+                  "Triton vs PyTorch int8"], rows)
+    (b_bf, p_bf), (b_tr, p_tr) = gemv_bandwidth(d, ms[0], "bf16"), gemv_bandwidth(d, ms[0], "triton")
+    calls = sum(r["per_step"] for r in d["rows"] if r["m"] == ms[0])
+    return text + (
+        f"\n\nTime for every matmul of one decode step (the {calls} Linears), replayed in CUDA graphs over weight "
+        f"copies larger than the L2 cache, so each call reads its weight from memory. At {ms[0]} row, a step must "
+        f"read {n(b_bf / 2**20)} MiB of bf16 weights or {n(b_tr / 2**20)} MiB of int8 weights and scales: bf16 "
+        f"runs at {p_bf:.0f}% of the measured {d['bandwidth'] / 1e9:.0f} GB/s, the Triton kernel at {p_tr:.0f}%. "
+        "PyTorch int8 casts each weight back to bf16 first, so it moves more bytes than bf16, not fewer.")
+
+
+def rmsnorm_section(d: dict) -> str:
+    batches = sorted({r["batch"] for r in d["rows"]})
+    rows = []
+    for b in batches:
+        t, k = need(d["rows"], "rmsnorm", variant="torch", batch=b), need(d["rows"], "rmsnorm", variant="triton", batch=b)
+        rows.append([str(b), n(t["decode_tok_s"]), n(k["decode_tok_s"]), _ratio(k["decode_tok_s"] / t["decode_tok_s"]),
+                     f"{t['ms_per_step']:.2f} → {k['ms_per_step']:.2f}"])
+    text = table(["Batch", "PyTorch norm tok/s", "Triton norm tok/s", "Triton vs PyTorch", "ms per step"], rows)
+    p, tf, alone = d["profile_b1"], d["teacher_forced"], d["kernels_alone"]
+    fused = sorted((r for r in alone if r["op"] == "add+norm"), key=lambda r: r["rows"])
+    if not fused or fused[0]["rows"] != 1:
+        raise MissingResult("rmsnorm alone: no add+norm row with 1 row")
+    one, big = fused[0], fused[-1]
+    return text + (
+        f"\n\nKV cache + CUDA graphs, the two norms alternating at each batch size. At batch 1 a decode step runs "
+        f"{n(p['torch']['gpu_ops_per_step'])} GPU ops with the PyTorch norm and {n(p['triton']['gpu_ops_per_step'])} "
+        f"with the Triton kernels (GPU busy {p['torch']['gpu_busy_ms_per_step']:.2f} → "
+        f"{p['triton']['gpu_busy_ms_per_step']:.2f} ms per step). Residual add + norm alone: one row takes "
+        f"{one['torch_us']:.1f} µs in PyTorch and {one['triton_us']:.1f} µs in Triton; {n(big['rows'])} rows run at "
+        f"{big['triton_pct_bw']:.0f}% of the measured {d['bandwidth'] / 1e9:.0f} GB/s. Greedy rows identical to the "
+        f"PyTorch norm: {d['greedy_rows_identical']}/{d['n_greedy_rows']}; teacher-forced argmax agreement "
+        f"{tf['argmax_agree']:.2%}.")
+
+
+def serving_baseline(rows: list[dict], r: dict) -> tuple[dict, str] | None:
+    """Static has none; continuous vs static at the same size; graphs vs continuous; paged vs 8 contiguous slots."""
+    e, s = r["engine"], r["size"]
+    if e == "continuous":
+        return row(rows, engine="static", size=s), "static"
+    if e == "cont+graph":
+        return row(rows, engine="continuous", size=s), "continuous"
+    if e.startswith("paged/"):
+        return row(rows, engine="continuous", size=8), "continuous 8 (same KV memory)"
+    return None
+
+
+def serving_table(d: dict) -> str:
+    rows = []
+    for r in d["rows"]:
+        base = serving_baseline(d["rows"], r)
+        vs = f"{r['tok_s'] / base[0]['tok_s']:.2f}× {base[1]}" if base and base[0] else "—"
+        alloc = r.get("kv_alloc_used_mean")
+        c = r.get("correctness") or {}
+        rows.append([r["engine"], str(r["size"]), n(r["tok_s"]), vs,
+                     f"{r['latency_mean_s']:.2f} / {r['latency_p95_s']:.2f}",
+                     f"{r['ttft_mean_s']:.2f} / {r['ttft_p95_s']:.2f}",
+                     f"{r['kv_reserved_mib']:.0f}", "—" if alloc is None else f"{alloc:.0%}",
+                     str(r["peak_active"] or "—"), str(r["preemptions"]),
+                     f"{c['identical']}/{c['of']}" if c else "—"])
+    w = d["workload"]
+    text = table(["Engine", "Slots", "tok/s", "vs", "Latency mean / p95 (s)", "TTFT mean / p95 (s)",
+                  "KV MiB", "KV claimed in use", "Peak active", "Preempt", "Same as solo"], rows)
+    reps = (d.get("provenance") or {}).get("repeats")
+    text += (f"\n\n{w['requests']} requests, all queued at t=0 (seed {w['seed']}); "
+             f"{'median-wall run of ' + str(reps) if reps else 'median run'}. \"Same as solo\": requests whose "
+             "tokens equal solo `generate()`; bf16 near-ties flip some, so it is reported, not asserted.")
+    return text
+
+
+def logits_section(d: dict, torchnorm: dict | None = None) -> str:
+    """`torchnorm`: the same run with our RMSNorm forced to PyTorch (EXTRAS); adds its rows, labelled by norm."""
+    rows = []
+    runs = [(d, "")] if torchnorm is None else [
+        (d, f", {(d.get('provenance') or {}).get('norm_backend', '?')} norm"), (torchnorm, ", torch norm")]
+    for (r, label), phase in [(x, p) for x in runs for p in ("prefill", "decode")]:
+        s = r["stats"][phase]
+        flips = s["flips"]["ours_bf16"]
+        rows.append([phase + label, str(s["ours_vs_hf_bf16"]["positions"]),
+                     "yes" if s["bit_identical_to_hf_bf16"] else "no",
+                     f"{s['ours_vs_hf_bf16']['max_abs']:.3g}",
+                     f"{s['ours_bf16']['top1']:.2%}", f"{s['hf_bf16']['top1']:.2%}",
+                     ", ".join(f"{f:.3f}" for f in flips) or "none"])
+    text = table(["Phase", "Positions", "Ours == HF bf16 (`torch.equal`)", "Max abs Δ, ours vs HF",
+                  "Top-1 vs fp32: ours", "Top-1 vs fp32: HF", "Our margin at each fp32 flip"], rows)
+    div = d["greedy_divergence"]
+    same = sum(g["ours_vs_hf_bf16"] is None for g in div)
+    at = lambda i: "never" if i is None else f"token {i}"                    # noqa: E731
+    text += (f"\n\nGreedy, {len(div)} prompts × {d['new_tokens']} tokens: ours equals HF bf16 on {same}/{len(div)}. "
+             "First divergence from the fp32 greedy text, ours / HF: "
+             + "; ".join(f"{at(g['ours_vs_fp32'])} / {at(g['hf_bf16_vs_fp32'])}" for g in div) + ".")
+    return text
+
+
+def profile_table(profiles: list[tuple[str, dict]]) -> str:
+    rows = []
+    for label, d in profiles:
+        for mode in ("eager", "graph"):
+            m = d["modes"][mode]
+            rows.append([label, mode, f"{m['ms_per_step']:.2f}", f"{m['host_issue_ms_per_step']:.2f}",
+                         f"{m['gpu_busy_ms_per_step']:.2f}", f"{m['gpu_idle_frac']:.0%}",
+                         n(m["decode_tok_s"]), f"{m['pct_of_ceiling']:.0f}%"])
+    return table(["Model (B=1)", "Decode step", "ms / step", "Host issue ms", "GPU busy ms", "GPU idle",
+                  "tok/s", "% of ceiling"], rows)
+
+
+# ---------------------------------------------------------------- pages
+
+
+def run_header(data: dict) -> str:
+    m = data["manifest"]
+    rec = m.get("benches") or {}
+    # the Run line describes the main run only; --only reruns get their own line, grouped by session
+    provs = [data[b.name].get("provenance") or {} for b in BENCHES if not (rec.get(b.name) or {}).get("rerun")]
+    commits = sorted({(p.get("git") or {}).get("commit") or "?" for p in provs})
+    dirty = any((p.get("git") or {}).get("dirty") for p in provs)
+    first = next((p for p in provs if p.get("gpu_name")), {})
+    times = sorted(p["time_utc"] for p in provs if p.get("time_utc"))
+    commit = (f"commit `{commits[0][:7]}`" if len(commits) == 1 else
+              f"⚠ {len(commits)} commits ({', '.join(c[:7] for c in commits)})")
+    lines = [
+        f"- **Run:** {times[0] if times else '?'} → {m.get('finished_utc', times[-1] if times else '?')} UTC, "
+        f"{commit}{' with uncommitted changes' if dirty else ''}, NitroSense **{m.get('power_mode', 'not recorded')}**.",
+        f"- **Machine:** {first.get('gpu_name', '?')}, driver {first.get('driver', '?')}, "
+        f"torch {first.get('torch', '?')}, CUDA {first.get('cuda', '?')}.",
+    ]
+    sessions: dict[tuple, list[str]] = {}
+    for n, r in sorted(rec.items(), key=lambda kv: str(kv[1].get("finished_utc", ""))):
+        if r.get("rerun"):
+            fin = str(r.get("finished_utc", "?"))
+            key = (fin[:10], ((r.get("git") or {}).get("commit") or "?")[:7], r.get("power_mode") or "not recorded")
+            sessions.setdefault(key, []).append(
+                f"{n} (finished {fin[11:16]} UTC, started at {r.get('start_temp_c', '?')} °C)")
+    if sessions:
+        lines.append("- **Rerun later with `--only`:** " + "; ".join(
+            f"{day}, commit `{c}`, NitroSense **{mode}**: {', '.join(names)}"
+            for (day, c, mode), names in sessions.items()) + ". Every other benchmark is from the run above.")
+    return "\n".join(lines)
+
+
+def render_results(data: dict, run_dir: str = "runs/bench/v1b") -> str:
+    tl_gen = [(nm, lab) for nm, lab in TL_ENGINES]
+    w = warnings(data)
+    parts = [
+        "# Results",
+        "",
+        f"<!-- Generated by scripts/bench_all.py from {run_dir}/. Do not edit by hand: change the script "
+        "(src/tinyllm/results.py) and re-render. -->",
+        "",
+        "Every number on this page comes from **one** run of `scripts/bench_all.py --run`: one commit, "
+        "one session, one GPU job at a time, with a cool-down before each benchmark"
+        + (" (a benchmark rerun later with `--only` is listed below)." if any(
+            r.get("rerun") for r in (data["manifest"].get("benches") or {}).values()) else "."),
+        "",
+        run_header(data),
+        "- **How to read it:** ratios between rows of one table are the reliable result; they share a "
+        "session and its heat. Absolute tok/s on this laptop moves 10–20% between sessions with "
+        "temperature, so each table ends with its heat line. ⚠ marks a hardware slowdown, or a thermal "
+        f"slowdown with the busy clock median below {n(MIN_CLOCK_MHZ)} MHz.",
+        "- **Elsewhere:** dated results and the reasoning behind each step are in "
+        "[INFERENCE.md](INFERENCE.md); ours vs vLLM vs llama.cpp under load (4.8a, rerun in 4.10c) is in "
+        "[SERVING.md](SERVING.md).",
+        "- **Reproduce:** `uv run python scripts/bench_all.py` re-renders this page and the README "
+        "tables from the saved JSON in seconds, with no GPU. `--run --power-mode <mode>` reruns "
+        "everything (~1.5–2 h); `--only <bench>` reruns one benchmark.",
+    ]
+    if w:
+        parts += ["", "**⚠ Throttling warnings:**", *[f"- {x}" for x in w]]
+    hf = data["tl-hf"]
+    parts += [
+        "", "## TinyLlama 1.1B (bf16)", "",
+        "### Decode at batch 1", "",
+        f"128-token prompt, 256 new tokens, {repeats_note(hf)}. Ceiling: {ceiling_note(hf)}.", "",
+        tl_b1_table(data), "", heat_line(data, [n_ for n_, _ in tl_gen], ["hf", "naive", "kv", "kv-graph"]),
+        "", img("img/decode_b1.svg", b1_alt(data)),
+        "", "### Decode by batch size and prompt length", "",
+        decode_grid(data, tl_gen), "", img("img/batch_scaling.svg", scaling_alt(data)), "",
+        "", "### INT8 weights (per-row scales)", "",
+        "Two ways to run an int8 matmul: `int8` casts the weight back to bf16 on every call (plain PyTorch); "
+        "`int8-triton` is our Triton kernel, which reads the int8 weight and dequantizes in registers (4.10b).",
+        "", "#### Eager decode (KV cache, no graphs)", "",
+        int8_section(data["tl-int8"]), "", heat_line(data, ["tl-int8"]),
+        "", "#### With CUDA graphs: PyTorch path vs Triton kernel", "",
+        "Each variant is timed with only its own weights on the GPU.", "",
+        int8_section(data["tl-int8-graph"]), "", heat_line(data, ["tl-int8-graph"]),
+        "", "#### The matmuls alone", "",
+        gemv_table(data["tl-int8-gemv"]), "", heat_line(data, ["tl-int8-gemv"]),
+        "", "### Fused RMSNorm + residual add (Triton, 4.10a)", "",
+        rmsnorm_section(data["tl-rmsnorm"]), "", heat_line(data, ["tl-rmsnorm"]),
+        "", "### Serving: static vs continuous batching, with and without CUDA graphs", "",
+        serving_table(data["tl-serving"]), "", heat_line(data, ["tl-serving"]),
+        "", "### Logits vs HF Transformers", "",
+        "Ours (bf16) and HF Transformers (bf16, eager) against an fp32 HF reference, on the same prompts.",
+        "", logits_section(data["tl-logits"], data.get("tl-logits-torchnorm")), "", heat_line(data, ["tl-logits"]),
+        "", "## Our 30M model (bf16 weights)", "",
+        "### Decode by batch size and prompt length", "",
+        f"256 new tokens, {repeats_note(data['30m-kv'])}. Ceiling: {ceiling_note(data['30m-kv'])}.", "",
+        decode_grid(data, [("30m-naive", "Naive"), ("30m-kv", "KV cache"), ("30m-kv-graph", "KV + CUDA graphs")]),
+        "", heat_line(data, ["30m-naive", "30m-kv", "30m-kv-graph"], ["naive", "kv", "kv-graph"]),
+        "", "### Serving: static, continuous, CUDA graphs, paged KV", "",
+        "Paged rows: `paged/<block size>`, all within the 8-slot contiguous cache's KV memory.", "",
+        serving_table(data["30m-serving"]), "", heat_line(data, ["30m-serving"]),
+        "", img("img/serving.svg", serving_alt(data)),
+        "", "### INT8 weights", "",
+        int8_section(data["30m-int8"]), "", heat_line(data, ["30m-int8"]),
+        "", "## Decode-step profile (batch 1, 128-token context)", "",
+        profile_table([("TinyLlama 1.1B", data["tl-profile"]), ("30M", data["30m-profile"])]), "",
+        heat_line(data, ["tl-profile", "30m-profile"], ["TinyLlama", "30M"]),
+        "", img("img/decode_profile.svg", profile_alt(data)),
+    ]
+    return re.sub(r"\n{3,}", "\n\n", "\n".join(parts)) + "\n"
+
+
+# ---------------------------------------------------------------- README blocks
+
+
+def _ratio(x: float, d: int = 2) -> str:
+    return f"{x:.{d}f}×"
+
+
+# ---------------------------------------------------------------- plot embeds (6.2, D-030)
+# The SVGs come from tinyllm.plots via bench_all.py; the alt text carries the plot's numbers.
+
+
+def img(path: str, alt: str) -> str:
+    return f"![{alt}]({path})"
+
+
+def _tl(data: dict, nm: str, b: int) -> dict:
+    return need(data[nm]["rows"], nm, batch=b, prompt_len=128)
+
+
+def b1_alt(data: dict) -> str:
+    hf = data["tl-hf"]
+    return ("Bar chart, TinyLlama 1.1B decode at batch 1: " + "; ".join(
+        f"{lab.split(' (')[0].replace('`', '')} {n(_tl(data, nm, 1)['decode_tok_s'])} tok/s" for nm, lab in TL_ENGINES)
+        + f"; bandwidth ceiling {n(hf['bandwidth_bytes_s'] / hf['weight_bytes'])} tok/s")
+
+
+def scaling_alt(data: dict) -> str:
+    bs = sorted({r["batch"] for r in data["tl-kv-graph"]["rows"]} & {r["batch"] for r in data["tl-hf"]["rows"]})
+    per = lambda nm: ", ".join(n(_tl(data, nm, b)["decode_tok_s"]) for b in bs)   # noqa: E731
+    return (f"Line chart, decode tok/s at batch {', '.join(map(str, bs))}: ours with CUDA graphs {per('tl-kv-graph')}; "
+            f"HF Transformers {per('tl-hf')}; log scales, with the bandwidth ceiling times batch size")
+
+
+def serving_alt(data: dict) -> str:
+    tls, ms = data["tl-serving"]["rows"], data["30m-serving"]["rows"]
+    tok = lambda rows, nm, **kv: n(need(rows, nm, **kv)["tok_s"])                 # noqa: E731
+    pg = need(ms, "30m-serving", engine=PAGED_ROW[0], size=PAGED_ROW[1])
+    return ("Bar chart, serving tok/s. TinyLlama, 8 slots: static " + tok(tls, "tl-serving", engine="static", size=8)
+            + ", continuous " + tok(tls, "tl-serving", engine="continuous", size=8)
+            + ", continuous with CUDA graphs " + tok(tls, "tl-serving", engine="cont+graph", size=8)
+            + ". 30M model, same KV memory: contiguous 8 sequences " + tok(ms, "30m-serving", engine="continuous", size=8)
+            + f", paged {pg['peak_active']} sequences {n(pg['tok_s'])}")
+
+
+def profile_alt(data: dict) -> str:
+    m = data["tl-profile"]["modes"]
+    return ("Bar chart, share of a TinyLlama decode step the GPU is idle: "
+            f"eager {m['eager']['gpu_idle_frac']:.0%} ({m['eager']['ms_per_step']:.1f} ms per step), "
+            f"CUDA graph {m['graph']['gpu_idle_frac']:.0%} ({m['graph']['ms_per_step']:.1f} ms per step)")
+
+
+def readme_blocks(data: dict) -> dict[str, str]:
+    g = {nm: data[nm]["rows"] for nm, _ in TL_ENGINES}
+    at = lambda nm, b: need(g[nm], nm, batch=b, prompt_len=128)              # noqa: E731
+    hf1, kv1, kvg1 = at("tl-hf", 1), at("tl-kv", 1), at("tl-kv-graph", 1)
+    tls, ms = data["tl-serving"]["rows"], data["30m-serving"]["rows"]
+    st8 = need(tls, "tl-serving", engine="static", size=8)
+    co8 = need(tls, "tl-serving", engine="continuous", size=8)
+    gr8 = need(tls, "tl-serving", engine="cont+graph", size=8)
+    c8 = need(ms, "30m-serving", engine="continuous", size=8)
+    pg = need(ms, "30m-serving", engine=PAGED_ROW[0], size=PAGED_ROW[1])
+    lg = data["tl-logits"]["stats"]
+    identical = lg["prefill"]["bit_identical_to_hf_bf16"] and lg["decode"]["bit_identical_to_hf_bf16"]
+    max_d = max(lg[p]["ours_vs_hf_bf16"]["max_abs"] for p in ("prefill", "decode"))
+    positions = sum(lg[p]["ours_vs_hf_bf16"]["positions"] for p in ("prefill", "decode"))
+    flips = [f for p in ("prefill", "decode") for f in lg[p]["flips"]["ours_bf16"]]
+    top1 = {p: lg[p]["ours_vs_hf_bf16"]["top1"] for p in ("prefill", "decode")}
+    tn = (data.get("tl-logits-torchnorm") or {}).get("stats")
+    torch_identical = bool(tn) and all(tn[p]["bit_identical_to_hf_bf16"] for p in ("prefill", "decode"))
+    gap_cause = (" The difference is our Triton RMSNorm (D-032): with the PyTorch norm the logits are "
+                 "bit-identical to HF's." if torch_identical else "")
+
+    logits_hl = ("**Our engine's logits are bit-identical to HF Transformers** in bf16 on TinyLlama, "
+                 "prefill and cached decode." if identical else
+                 f"**Our engine's bf16 logits match HF Transformers to within {max_d:.2g}** on TinyLlama "
+                 f"(same top token at {top1['prefill']:.1%} of prefill and {top1['decode']:.1%} of cached-decode "
+                 "positions)." + gap_cause)
+    highlights = "\n".join([
+        img("docs/img/decode_b1.svg", b1_alt(data)),
+        "",
+        f"- **TinyLlama 1.1B decodes at {n(kvg1['decode_tok_s'])} tok/s at batch 1: "
+        f"{kvg1['pct_of_ceiling']:.0f}% of the bandwidth ceiling, "
+        f"{_ratio(kvg1['decode_tok_s'] / hf1['decode_tok_s'], 1)} HF `generate`** (KV cache + CUDA graphs).",
+        f"- {logits_hl}",
+        f"- **Serving:** continuous batching gives {_ratio(co8['tok_s'] / st8['tok_s'])} static batching. "
+        f"A paged KV cache fits {pg['peak_active']} sequences where contiguous slots fit 8, for "
+        f"{_ratio(pg['tok_s'] / c8['tok_s'])} the throughput in the same memory.",
+    ])
+
+    hf = data["tl-hf"]
+    prof = data["tl-profile"]["modes"]
+    kv32 = at("tl-kv", 32)
+    host_bound = prof["eager"]["gpu_idle_frac"] >= 0.5
+    i8 = data["tl-int8-graph"]
+    vs_bf16 = lambda v, b: (need(i8["rows"], "tl-int8-graph", variant=v, batch=b)["decode_tok_s"]   # noqa: E731
+                            / need(i8["rows"], "tl-int8-graph", variant="bf16", batch=b)["decode_tok_s"])
+    tri, cast = {b: vs_bf16("int8-triton", b) for b in (1, 32)}, {b: vs_bf16("int8", b) for b in (1, 32)}
+    int8_text = (f"- **INT8 weights:** val loss {i8['val_loss']['int8-triton'] - i8['val_loss']['bf16']:+.4f}, "
+                 f"{i8['teacher_forced']['int8-triton']['argmax_agree']:.1%} argmax agreement with bf16. Our "
+                 "Triton kernel reads the int8 weight and dequantizes inside the matmul; with CUDA graphs, decode "
+                 f"is {_ratio(tri[1])} bf16 at batch 1 and {_ratio(tri[32])} at batch 32. Casting the weight back "
+                 f"to bf16 on every call (plain PyTorch) gives {_ratio(cast[1])} and {_ratio(cast[32])}")
+    int8_text += (": it moves more bytes than bf16, not fewer." if cast[1] < 1 and cast[32] < 1 else ".")
+    eager_text = (f"- **Eager decode is {'host-bound' if host_bound else 'slow'}.** At batch 1, HF and our eager "
+                  f"engine take {hf1['ms_per_step']:.1f} and {kv1['ms_per_step']:.1f} ms per step "
+                  f"(ours {kv32['ms_per_step']:.1f} ms at batch 32). The profiler shows the GPU idle "
+                  f"{prof['eager']['gpu_idle_frac']:.0%} of an eager step"
+                  + (", waiting for kernel launches" if host_bound else "")
+                  + f". Capturing the decode step as a CUDA graph removes that overhead; the GPU is then idle "
+                  f"{prof['graph']['gpu_idle_frac']:.0%} of the step.")
+    corr = (f"- **Correctness:** bf16 logits equal HF's bf16 logits exactly (`torch.equal`) at all {positions} "
+            "test positions." if identical else
+            f"- **Correctness:** bf16 logits differ from HF's bf16 logits by at most {max_d:.2g} over "
+            f"{positions} test positions." + gap_cause)
+    corr += (f" Every top-1 disagreement with fp32 is a near-tie (margin ≤ {max(flips):.3f} logits)." if flips
+             else " No top-1 disagreement with fp32.")
+    tinyllama = "\n".join([
+        f"Decode speed at batch 1 with a 128-token prompt, 256 new tokens, {repeats_note(hf)}. The ceiling is "
+        f"the measured bandwidth divided by the weight size: {ceiling_note(hf)}, because each decoded token "
+        "has to read every weight once.",
+        "",
+        tl_b1_table(data),
+        "",
+        eager_text,
+        f"- **Batching:** with graphs, {n(at('tl-kv-graph', 8)['decode_tok_s'])} tok/s at batch 8 and "
+        f"{n(at('tl-kv-graph', 32)['decode_tok_s'])} tok/s at batch 32 (HF: {n(at('tl-hf', 8)['decode_tok_s'])} "
+        f"and {n(at('tl-hf', 32)['decode_tok_s'])}).",
+        int8_text,
+        corr,
+        "",
+        img("docs/img/batch_scaling.svg", scaling_alt(data)),
+    ])
+
+    w = data["tl-serving"]["workload"]
+    reps = (data["tl-serving"].get("provenance") or {}).get("repeats")
+    serving = "\n".join([
+        f"{w['requests']} requests of mixed lengths ({n(st8['useful_tokens'])} tokens to generate), "
+        f"{'median of ' + str(reps) + ' runs' if reps else 'median run'}.",
+        "",
+        table(["Scheduler / cache", "Setting", "tok/s", "vs baseline"], [
+            ["Static batching (TinyLlama)", "8 slots", n(st8["tok_s"]), "—"],
+            ["Continuous batching (TinyLlama)", "8 slots", n(co8["tok_s"]), _ratio(co8["tok_s"] / st8["tok_s"])],
+            ["Continuous batching + CUDA graphs (TinyLlama)", "8 slots", n(gr8["tok_s"]),
+             f"{_ratio(gr8['tok_s'] / co8['tok_s'])} vs eager continuous"],
+            ["Contiguous KV slots (30M model)", f"{c8['kv_reserved_mib']:.0f} MiB of KV, 8 sequences",
+             n(c8["tok_s"]), "—"],
+            ["Paged KV cache (30M model)", f"same {pg['kv_reserved_mib']:.0f} MiB, {pg['peak_active']} sequences",
+             n(pg["tok_s"]), _ratio(pg["tok_s"] / c8["tok_s"])],
+        ]),
+        "",
+        "The paged cache stores KV in fixed-size blocks through a per-sequence block table, as in vLLM, so a "
+        "sequence holds only the blocks it has filled. The share of reserved KV holding real tokens goes from "
+        f"{c8['kv_alloc_used_mean']:.0%} to {pg['kv_alloc_used_mean']:.0%}.",
+    ])
+    return {"highlights": highlights, "tinyllama": tinyllama, "serving": serving}
+
+
+def replace_block(text: str, name: str, body: str) -> str:
+    """Replace what sits between <!-- bench:NAME --> and <!-- /bench:NAME -->; the rest stays byte-identical."""
+    start, end = f"<!-- bench:{name} -->", f"<!-- /bench:{name} -->"
+    if text.count(start) != 1 or text.count(end) != 1:
+        raise ValueError(f"expected exactly one {start} and one {end}")
+    i, j = text.index(start) + len(start), text.index(end)
+    if j < i:
+        raise ValueError(f"{end} comes before {start}")
+    return text[:i] + "\n" + body.strip("\n") + "\n" + text[j:]
+
+
+def render_readme(text: str, data: dict) -> str:
+    for name, body in readme_blocks(data).items():
+        text = replace_block(text, name, body)
+    return text
+
+
+def block_names(text: str) -> list[str]:
+    return re.findall(r"<!-- bench:([\w-]+) -->", text)

src/tinyllm/safetensors_io.py

@@ -0,0 +1,90 @@
+"""Minimal .safetensors reader and writer (step 4.9b, D-026). No dependency beyond torch.
+
+File layout (https://github.com/huggingface/safetensors#format):
+  8 bytes   little-endian u64 N = header size
+  N bytes   UTF-8 JSON: {name: {"dtype": "BF16", "shape": [...], "data_offsets": [begin, end]}, ...}
+            plus an optional "__metadata__": {str: str}. Padded with spaces to a multiple of 8.
+  rest      raw little-endian tensor bytes; offsets are relative to the end of the header
+
+Reading memory-maps the file, so a tensor is copied once, straight to its target device.
+"""
+
+from __future__ import annotations
+
+import json
+import mmap
+import struct
+from pathlib import Path
+
+import torch
+
+DTYPES = {
+    "BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32, "F64": torch.float64,
+    "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
+    "U8": torch.uint8, "BOOL": torch.bool,
+}
+NAMES = {v: k for k, v in DTYPES.items()}
+
+
+def read_header(path: str | Path) -> tuple[dict, int]:
+    """(header without __metadata__, byte offset where the data starts)."""
+    with open(path, "rb") as f:
+        (n,) = struct.unpack("<Q", f.read(8))
+        header = json.loads(f.read(n))
+    header.pop("__metadata__", None)
+    return header, 8 + n
+
+
+def load_file(path: str | Path, device: str | torch.device = "cpu", dtype: torch.dtype | None = None,
+              keys: set[str] | None = None) -> dict[str, torch.Tensor]:
+    """name -> tensor on `device`. dtype: cast floating tensors to it (integers are kept)."""
+    header, start = read_header(path)
+    out = {}
+    with open(path, "rb") as f:
+        # ACCESS_COPY: a writable (copy-on-write) view, so torch.frombuffer doesn't warn.
+        # Nothing is written back and pages are only copied if touched for writing.
+        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY)
+        try:
+            for name, info in header.items():
+                if keys is not None and name not in keys:
+                    continue
+                dt = DTYPES[info["dtype"]]
+                begin, end = info["data_offsets"]
+                shape = info["shape"]
+                numel = 1
+                for s in shape:
+                    numel *= s
+                if (end - begin) != numel * torch.empty((), dtype=dt).element_size():
+                    raise ValueError(f"{name}: {end - begin} bytes for shape {shape} {info['dtype']}")
+                if numel == 0:
+                    t = torch.empty(shape, dtype=dt)
+                else:
+                    t = torch.frombuffer(mm, dtype=dt, count=numel, offset=start + begin).view(shape)
+                target = dtype if (dtype is not None and dt.is_floating_point) else dt
+                out[name] = t.to(device=device, dtype=target, copy=True)   # copy: detach from the mmap
+                del t                                    # release the view, or mm.close() raises BufferError
+        finally:
+            mm.close()
+    return out
+
+
+def save_file(tensors: dict[str, torch.Tensor], path: str | Path, metadata: dict[str, str] | None = None) -> None:
+    """Write tensors (any device) in safetensors layout. Shared tensors are written once per name."""
+    header: dict = {}
+    blobs = []
+    offset = 0
+    for name in sorted(tensors):
+        t = tensors[name].detach().contiguous().cpu()
+        raw = t.reshape(-1).view(torch.uint8).numpy().tobytes() if t.numel() else b""
+        header[name] = {"dtype": NAMES[t.dtype], "shape": list(t.shape), "data_offsets": [offset, offset + len(raw)]}
+        blobs.append(raw)
+        offset += len(raw)
+    if metadata:
+        header["__metadata__"] = metadata
+    h = json.dumps(header, separators=(",", ":")).encode()
+    h += b" " * (-len(h) % 8)
+    with open(path, "wb") as f:
+        f.write(struct.pack("<Q", len(h)))
+        f.write(h)
+        for raw in blobs:
+            f.write(raw)

src/tinyllm/server.py

@@ -0,0 +1,1031 @@
+"""Step 4.13 step 2: OpenAI-compatible HTTP server over ContinuousEngine (D-025, D-028).
+
+Design: docs/SERVING.md.
+
+    client ──HTTP──▶ FastAPI handler ──▶ Router ──▶ EngineWorker thread: ContinuousEngine.step()
+       ◀── JSON / SSE ◀── per-request asyncio.Queue ◀── tokens (call_soon_threadsafe)
+
+- One thread owns the engine (it is not thread-safe). Handlers talk to it through an inbox.
+- POST /v1/completions: prompt as text or token ids, max_tokens, temperature, top_p, top_k,
+  ignore_eos, stream (SSE ending with `data: [DONE]`), finish_reason, usage.
+- Hardening:
+  - a bounded queue → 429 + Retry-After
+  - a per-client token-bucket rate limit → 429
+  - a per-request timeout → 504, slot freed
+  - client disconnect → abort, slot freed
+  - graceful shutdown
+  - an engine error → 500 for in-flight requests and not ready
+  - out of memory → 503 for in-flight requests, then recover
+- Observability (obs.py): a JSON log line per request, X-Request-ID, engine snapshots,
+  /stats, /metrics (Prometheus text format), --log-level debug with per-step timings,
+  slow-request timelines, and /debug/state (localhost only).
+- Pages (6.5b): GET / is the demo (story continuation, streamed), GET /status the status page
+  (charts and SLOs from GET /stats/history: the last hour per minute, the 7-day error budget).
+- Router with one variant now; X-Model-Version on every response. Release strategies (6.9)
+  add routing rules.
+
+Run the local GPU server (TinyLlama) or our 30M model:
+
+    uv run python -m tinyllm.server --model tinyllama --slots 8 --cuda-graph
+    uv run python -m tinyllm.server --ckpt checkpoints/main/best.pt --slots 8
+
+or our 30M model in HF format (6.5b), from a folder or from the Hub (the Docker image does this):
+
+    uv run python -m tinyllm.server --hf-model runs/export/tinystories-30m
+    uv run python -m tinyllm.server --hf-model CoolDip00/tiny-llm-tinystories-30m --device cpu
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import concurrent.futures
+import hashlib
+import json
+import logging
+import math
+import os
+import queue
+import re
+import secrets
+import subprocess
+import threading
+import time
+import uuid
+from contextlib import asynccontextmanager
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Protocol
+
+import torch
+from fastapi import FastAPI
+from fastapi import Request as HTTPRequest
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, StreamingResponse
+from pydantic import BaseModel, Field
+
+from tinyllm.engine import ContinuousEngine, Request, Sampling
+from tinyllm.model import norm_backend
+from tinyllm.obs import EngineSnapshot, Metrics, RequestRecord, hash_client, itl_stats, setup_logging
+
+log = logging.getLogger("tinyllm.server")
+
+
+class Tokenizer(Protocol):
+    eos_id: int | None
+
+    def encode(self, text: str) -> list[int]: ...
+    def decode(self, ids: list[int]) -> str: ...
+
+
+# -- engine thread ----------------------------------------------------------------------
+
+@dataclass
+class Faults:
+    """Test-only fault injection (D-028 amendment): forces the timeout / 500 / OOM paths."""
+    step_delay_s: float = 0.0           # sleep before every step (slow engine)
+    fail_at_step: int | None = None     # RuntimeError at step N (1-based): fatal, not ready
+    oom_at_step: int | None = None      # torch.OutOfMemoryError at step N: recoverable
+
+
+@dataclass
+class Item:
+    """What the engine thread sends a request's handler."""
+    kind: str                           # "token" or "error"
+    token: int = 0
+    t: float = 0.0                      # time.monotonic() when the token was produced
+    finish: str | None = None           # "stop" / "length" on the last token
+    queue_wait: float | None = None     # on the first token: arrival → start of its step
+    status: int = 200                   # errors: 400, 500, 503, 504
+    message: str = ""
+
+
+@dataclass
+class Job:
+    request_id: str
+    request: Request
+    loop: asyncio.AbstractEventLoop
+    out: asyncio.Queue
+    arrival: float                      # time.monotonic()
+    deadline: float
+    rid: int | None = None              # engine id, set by the engine thread
+    n_tokens: int = 0
+
+    def send(self, item: Item) -> None:
+        self.loop.call_soon_threadsafe(self.out.put_nowait, item)
+
+
+class EngineWorker:
+    """Owns one ContinuousEngine on its own thread. submit() / abort() / debug_state() are
+    safe to call from the event loop; everything else runs on the engine thread."""
+
+    def __init__(self, engine: ContinuousEngine, metrics: Metrics, max_queue: int,
+                 snapshot_s: float = 5.0, faults: Faults | None = None, name: str = "main"):
+        self.engine = engine
+        self.metrics = metrics
+        self.max_outstanding = engine.n_slots + max_queue
+        self.snapshot_s = snapshot_s
+        self.faults = faults or Faults()
+        self.name = name
+        self.accepting = True
+        self.failed: str | None = None
+        self.steps = 0
+        self.busy = 0
+        self._outstanding = 0
+        self._lock = threading.Lock()
+        self._inbox: queue.SimpleQueue = queue.SimpleQueue()
+        self._jobs: dict[int, Job] = {}            # engine rid → job (engine thread only)
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, name=f"engine-{name}", daemon=True)
+        self._tokens_since_snap = 0
+        self._last_snap = time.monotonic()
+
+    # -- event-loop side ------------------------------------------------------------
+    def start(self) -> None:
+        self._thread.start()
+
+    @property
+    def alive(self) -> bool:
+        return self._thread.is_alive()
+
+    def load(self) -> tuple[int, int]:
+        """(queued, busy) as last seen; a gauge, good enough for logs and admission."""
+        with self._lock:
+            return max(0, self._outstanding - self.busy), self.busy
+
+    def submit(self, job: Job) -> str | None:
+        """Queue a job. Returns None, or why it was refused: "failed", "draining", "full"."""
+        with self._lock:
+            if self.failed:
+                return "failed"
+            if not self.accepting:
+                return "draining"
+            if self._outstanding >= self.max_outstanding:
+                return "full"
+            self._outstanding += 1
+        self._inbox.put(("add", job))
+        return None
+
+    def abort(self, job: Job) -> None:
+        self._inbox.put(("abort", job))
+
+    async def debug_state(self) -> dict:
+        fut: concurrent.futures.Future = concurrent.futures.Future()
+        self._inbox.put(("debug", fut))
+        return await asyncio.wait_for(asyncio.wrap_future(fut), 5.0)
+
+    def stop(self, drain_s: float) -> None:
+        """Stop accepting, let in-flight requests finish (up to drain_s), then stop."""
+        with self._lock:
+            self.accepting = False
+        end = time.monotonic() + drain_s
+        while time.monotonic() < end and self.outstanding() and self.alive:
+            time.sleep(0.02)
+        self._stop.set()
+        self._thread.join(timeout=10)
+
+    def outstanding(self) -> int:
+        with self._lock:
+            return self._outstanding
+
+    # -- engine thread --------------------------------------------------------------
+    def _finish(self, job: Job) -> None:
+        if job.rid is not None:
+            self._jobs.pop(job.rid, None)
+            self.engine.results.pop(job.rid, None)
+        with self._lock:
+            self._outstanding -= 1
+
+    def _fail_all(self, status: int, message: str) -> None:
+        for job in list(self._jobs.values()):
+            if self.failed is None:
+                self.engine.abort(job.rid)
+            job.send(Item("error", status=status, message=message))
+            self._finish(job)
+
+    def _handle(self, cmd: str, arg) -> None:
+        if cmd == "add":
+            job: Job = arg
+            if self.failed or self._stop.is_set():
+                job.send(Item("error", status=503 if not self.failed else 500,
+                              message="engine failed" if self.failed else "shutting down"))
+                self._finish(job)
+                return
+            try:
+                job.rid = self.engine.add_request(job.request)
+            except ValueError as e:
+                job.send(Item("error", status=400, message=str(e)))
+                self._finish(job)
+                return
+            self._jobs[job.rid] = job
+        elif cmd == "abort":
+            job = arg
+            if job.rid is not None and self._jobs.get(job.rid) is job:
+                if self.failed is None:
+                    self.engine.abort(job.rid)
+                self._finish(job)
+        elif cmd == "debug":
+            arg.set_result(self._debug_state())
+
+    def _drain_inbox(self, block: bool) -> None:
+        try:
+            cmd, arg = self._inbox.get(timeout=0.05) if block else self._inbox.get_nowait()
+        except queue.Empty:
+            return
+        self._handle(cmd, arg)
+        while True:
+            try:
+                cmd, arg = self._inbox.get_nowait()
+            except queue.Empty:
+                return
+            self._handle(cmd, arg)
+
+    def _expire(self, now: float) -> None:
+        for job in [j for j in self._jobs.values() if now > j.deadline]:
+            self.engine.abort(job.rid)
+            job.send(Item("error", status=504, message="request timed out"))
+            self._finish(job)
+
+    def _step(self) -> None:
+        self.steps += 1
+        f = self.faults
+        if f.step_delay_s:
+            time.sleep(f.step_delay_s)
+        if f.fail_at_step == self.steps:
+            raise RuntimeError(f"injected fault at step {self.steps}")
+        if f.oom_at_step == self.steps:
+            raise torch.OutOfMemoryError(f"injected OOM at step {self.steps}")
+        t0 = time.monotonic()
+        events = self.engine.step()
+        t1 = time.monotonic()
+        for e in events:
+            job = self._jobs.get(e.rid)
+            if job is None:
+                continue
+            job.n_tokens += 1
+            finish = None
+            if e.done:
+                eos = job.request.eos_id
+                finish = "stop" if eos is not None and e.token == eos else "length"
+            job.send(Item("token", e.token, t1, finish,
+                          queue_wait=t0 - job.arrival if job.n_tokens == 1 else None))
+            if e.done:
+                self._finish(job)
+        self._tokens_since_snap += len(events)
+        queued, busy = self.engine.load()
+        with self._lock:
+            self.busy = busy
+        if log.isEnabledFor(logging.DEBUG):
+            log.debug("step", extra={"fields": {"step": self.steps, "step_ms": round(1e3 * (t1 - t0), 3),
+                                                "tokens": len(events), "queued": queued, "busy": busy}})
+
+    def _snapshot(self, now: float) -> None:
+        dt = max(now - self._last_snap, 1e-9)
+        queued, busy = self.engine.load()
+        snap = EngineSnapshot(time.time(), queued, busy, self.engine.n_slots,
+                              round(self._tokens_since_snap / dt, 2), round(self.engine.kv_used_frac(), 4),
+                              ram_mb=_rss_mb(), vram_mb=_vram_mb(self.engine.device))
+        self.metrics.snapshot(snap)
+        if busy or queued or self._tokens_since_snap:
+            log.info("engine", extra={"fields": {"variant": self.name, **asdict(snap)}})
+        self._tokens_since_snap, self._last_snap = 0, now
+
+    def _run(self) -> None:
+        self._snapshot(time.monotonic())
+        while not self._stop.is_set():
+            idle = self.failed is not None or not self.engine.has_work()
+            self._drain_inbox(block=idle)
+            now = time.monotonic()
+            if now - self._last_snap >= self.snapshot_s and self.failed is None:
+                self._snapshot(now)
+            if self.failed is not None:
+                continue
+            self._expire(now)
+            if not self.engine.has_work():
+                continue
+            try:
+                self._step()
+            except torch.OutOfMemoryError as e:
+                log.error("out of memory: aborting in-flight requests", exc_info=True,
+                          extra={"fields": {"variant": self.name, "in_flight": len(self._jobs)}})
+                self._fail_all(503, f"out of memory, retry later ({e})")
+                if self.engine.device.type == "cuda":
+                    torch.cuda.empty_cache()
+                if self.engine.has_work():                       # could not get back to a clean state
+                    self.failed = "out of memory, engine not clean"
+            except Exception as e:                              # engine state unknown: stop serving
+                self.failed = f"{type(e).__name__}: {e}"
+                log.error("engine step failed: not ready", exc_info=True,
+                          extra={"fields": {"variant": self.name, "in_flight": len(self._jobs)}})
+                self._fail_all(500, "engine error")
+                with self._lock:
+                    self.busy = 0
+        # stopped: answer whatever is left
+        self._fail_all(503, "server shutting down")
+        while True:
+            try:
+                cmd, arg = self._inbox.get_nowait()
+            except queue.Empty:
+                break
+            self._handle(cmd, arg)
+
+    def _debug_state(self) -> dict:
+        now = time.monotonic()
+        queued, busy = self.engine.load()
+        return {"variant": self.name, "steps": self.steps, "failed": self.failed, "accepting": self.accepting,
+                "queued": queued, "busy": busy, "n_slots": self.engine.n_slots,
+                "kv_used_frac": round(self.engine.kv_used_frac(), 4),
+                "jobs": [{"request_id": j.request_id, "rid": j.rid, "age_s": round(now - j.arrival, 3),
+                          "tokens": j.n_tokens, "max_tokens": j.request.max_new_tokens,
+                          "deadline_in_s": round(j.deadline - now, 3)} for j in self._jobs.values()]}
+
+
+def _rss_mb() -> float | None:
+    try:
+        for line in Path("/proc/self/status").read_text().splitlines():
+            if line.startswith("VmRSS:"):
+                return round(int(line.split()[1]) / 1024, 1)
+    except OSError:
+        pass
+    return None
+
+
+def _vram_mb(device: torch.device) -> float | None:
+    if device.type != "cuda":
+        return None
+    return round(torch.cuda.memory_allocated(device) / 2**20, 1)
+
+
+# -- routing, limits --------------------------------------------------------------------
+
+@dataclass
+class Variant:
+    name: str
+    version: str                        # checkpoint hash + git commit
+    worker: EngineWorker
+    tokenizer: Tokenizer
+    max_len: int
+    vocab_size: int
+
+
+class Router:
+    """Picks the variant for a request. One variant now; 6.9 adds blue-green / canary /
+    shadow / A/B rules here, so handlers never change."""
+
+    def __init__(self, variants: list[Variant]):
+        self.variants = {v.name: v for v in variants}
+        self.primary = variants[0]
+
+    def pick(self, client: str) -> Variant:
+        return self.primary
+
+
+class RateLimiter:
+    """Token bucket per client: `per_min` requests per minute, bursts up to `burst`.
+    Called from the event loop only. per_min <= 0 disables it."""
+
+    def __init__(self, per_min: float, burst: int):
+        self.rate = per_min / 60.0
+        self.burst = max(1, burst)
+        self._b: dict[str, tuple[float, float]] = {}
+
+    def check(self, client: str, now: float) -> float:
+        """0 if allowed (and counted), else seconds until the next request is allowed."""
+        if self.rate <= 0:
+            return 0.0
+        tokens, last = self._b.get(client, (float(self.burst), now))
+        tokens = min(float(self.burst), tokens + (now - last) * self.rate)
+        if tokens >= 1.0:
+            self._b[client] = (tokens - 1.0, now)
+            if len(self._b) > 10_000:                     # forget clients whose bucket is full again
+                self._b = {k: v for k, v in self._b.items() if v[0] + (now - v[1]) * self.rate < self.burst}
+            return 0.0
+        self._b[client] = (tokens, now)
+        return (1.0 - tokens) / self.rate
+
+
+# -- HTTP layer -------------------------------------------------------------------------
+
+class CompletionRequest(BaseModel):
+    """OpenAI /v1/completions subset, plus vLLM's top_k / ignore_eos / return_token_ids."""
+    model: str | None = None
+    prompt: str | list[int]
+    max_tokens: int = Field(16, ge=1)
+    temperature: float = Field(1.0, ge=0.0)
+    top_p: float | None = Field(None, gt=0.0, le=1.0)
+    top_k: int | None = Field(None, ge=1)
+    stream: bool = False
+    ignore_eos: bool = False
+    return_token_ids: bool = False
+
+
+@dataclass
+class ServerConfig:
+    max_queue: int = 32                  # waiting requests beyond the slots; more → 429
+    timeout_s: float = 60.0              # per request, arrival → last token; then 504
+    rate_per_min: float = 0.0            # per client; 0 = off (the Space sets it)
+    rate_burst: int = 10
+    slow_s: float = 10.0                 # requests slower than this log their timeline (warning)
+    trust_proxy: bool = False            # take the client IP from X-Forwarded-For (HF Space)
+    proxy_hops: int = 1                  # with trust_proxy: trusted proxies in front, each appending to it
+    salt: str | None = None              # client-hash salt; None = resolve_salt()'s order
+    snapshot_s: float = 5.0
+    drain_s: float = 30.0
+
+
+_RID = re.compile(r"^[A-Za-z0-9._-]{1,64}$")
+
+
+class HeadersMiddleware:
+    """Pure ASGI (keeps streaming and disconnects intact): every response gets X-Request-ID
+    (the client's, if it sent a sane one) and X-Model-Version."""
+
+    def __init__(self, app, default_version: str):
+        self.app = app
+        self.default_version = default_version
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+        incoming = dict(scope.get("headers") or []).get(b"x-request-id", b"").decode("latin-1")
+        rid = incoming if _RID.match(incoming) else uuid.uuid4().hex[:16]
+        state = scope.setdefault("state", {})
+        state["request_id"] = rid
+
+        async def send_with_headers(msg):
+            if msg["type"] == "http.response.start":
+                version = state.get("model_version", self.default_version)
+                msg = {**msg, "headers": [*msg.get("headers", []), (b"x-request-id", rid.encode()),
+                                          (b"x-model-version", version.encode())]}
+            await send(msg)
+
+        await self.app(scope, receive, send_with_headers)
+
+
+class _ClientGone(Exception):
+    pass
+
+
+WEB = Path(__file__).resolve().parent / "web"
+
+
+def _page(name: str) -> str:
+    """The demo and status pages (6.5b): static HTML + JS in src/tinyllm/web, read per request
+    (small) so an edit shows on reload. They fetch the JSON endpoints with relative URLs."""
+    return (WEB / name).read_text(encoding="utf-8")
+
+
+def decode_completion(tok: Tokenizer, ids: list[int], context: list[int]) -> str:
+    """Completion text as it reads after the prompt. Decoded behind the prompt's last token(s),
+    because a SentencePiece decoder (TinyLlama) drops a leading space at the start of the text:
+    "…named Max." + " Max was" must not come back as "Max was" (byte-level BPE keeps it anyway).
+    Falls back to a plain decode if the context's text isn't a prefix (a partial UTF-8 byte)."""
+    if context:
+        head, text = tok.decode(context), tok.decode(context + ids)
+        if text.startswith(head):
+            return text[len(head):]
+    return tok.decode(ids)
+
+
+class _Detokenizer:
+    """Incremental text for streaming: decode everything so far, emit the new suffix; hold
+    back while the text ends in a partial UTF-8 character."""
+
+    def __init__(self, tok: Tokenizer, context: list[int] | None = None):
+        self.tok, self.ids, self.sent, self.context = tok, [], "", list(context or [])
+
+    def add(self, token: int) -> str:
+        self.ids.append(token)
+        return self._delta(final=False)
+
+    def flush(self) -> str:
+        return self._delta(final=True)
+
+    def _delta(self, final: bool) -> str:
+        text = decode_completion(self.tok, self.ids, self.context)
+        if (not final and text.endswith("�")) or not text.startswith(self.sent):
+            return ""
+        d, self.sent = text[len(self.sent):], text
+        return d
+
+
+def forwarded_ip(xff: str, hops: int = 1) -> str | None:
+    """The client's address in an X-Forwarded-For header behind `hops` trusted proxies.
+
+    Each proxy appends the address the request came from, so only the last `hops` entries
+    are trustworthy; anything left of them is whatever the client sent (spoofable, so never
+    the leftmost). With fewer entries than hops, all came from proxies: take the leftmost.
+    """
+    parts = [p.strip() for p in xff.split(",") if p.strip()]
+    if not parts:
+        return None
+    return parts[-hops] if len(parts) >= hops else parts[0]
+
+
+def resolve_salt(salt: str | None = None) -> tuple[str, str]:
+    """(salt, where it came from) for the client hash. The source is logged, never the salt.
+
+    Order: an explicit salt; $TINYLLM_LOG_SALT (a Space secret, the same value in every
+    replica Space, D-031); a hash of $SPACE_SIGNING_SECRET (HF's per-Space secret, stable
+    across restarts, but different in each Space); else random per start.
+    """
+    if salt:
+        return salt, "config"
+    if s := os.environ.get("TINYLLM_LOG_SALT"):
+        return s, "TINYLLM_LOG_SALT"
+    if s := os.environ.get("SPACE_SIGNING_SECRET"):
+        return hashlib.sha256(f"tinyllm-log-salt:{s}".encode()).hexdigest(), "SPACE_SIGNING_SECRET"
+    return secrets.token_hex(16), "random"
+
+
+def cpu_limit(cgroup: str = "/sys/fs/cgroup") -> int | None:
+    """CPUs this container may use, or None if no limit is known.
+
+    torch sizes its thread pool from the host's cores and ignores the container's CPU quota:
+    under `docker run --cpus 2` on a 16-core host it starts 8 threads, and decode ran ~5x
+    slower than with 2 (6.5b step 5, runs/bench/load/space-cpu/probe-threads-*). Order: the
+    cgroup v2 quota (`cpu.max`, "200000 100000" = 2 CPUs), cgroup v1, then HF's $CPU_CORES.
+    """
+    root = Path(cgroup)
+    try:
+        quota, period = (root / "cpu.max").read_text().split()[:2]
+        if quota != "max":
+            return max(1, math.ceil(int(quota) / int(period)))
+    except (OSError, ValueError):
+        pass
+    try:
+        quota = int((root / "cpu" / "cpu.cfs_quota_us").read_text())
+        period = int((root / "cpu" / "cpu.cfs_period_us").read_text())
+        if quota > 0 and period > 0:
+            return max(1, math.ceil(quota / period))
+    except (OSError, ValueError):
+        pass
+    try:
+        n = math.ceil(float(os.environ.get("CPU_CORES", "")))
+        return n if n > 0 else None
+    except ValueError:
+        return None
+
+
+class Server:
+    """State shared by the handlers (app.state.srv)."""
+
+    def __init__(self, variants: list[Variant], cfg: ServerConfig, metrics: Metrics):
+        self.router = Router(variants)
+        self.cfg = cfg
+        self.metrics = metrics
+        self.limiter = RateLimiter(cfg.rate_per_min, cfg.rate_burst)
+        self.salt, self.salt_source = resolve_salt(cfg.salt)
+        self.started = False
+
+    def client_id(self, request: HTTPRequest) -> str:
+        ip = request.client.host if request.client else "unknown"
+        if self.cfg.trust_proxy and (xff := request.headers.get("x-forwarded-for")):
+            ip = forwarded_ip(xff, self.cfg.proxy_hops) or ip
+        return hash_client(ip, self.salt)
+
+    def record(self, rec: RequestRecord, timeline: dict | None = None) -> None:
+        self.metrics.observe(rec)
+        fields = rec.to_dict()
+        if rec.latency_s is not None and rec.latency_s > self.cfg.slow_s:
+            log.warning("slow request", extra={"fields": {**fields, "timeline": timeline or {}}})
+        else:
+            log.info("request", extra={"fields": fields})
+
+    def error(self, rec: RequestRecord, status: int, message: str, headers: dict | None = None,
+              err_type: str = "invalid_request_error") -> JSONResponse:
+        rec.status, rec.error = status, message
+        rec.finish_reason = rec.finish_reason or ("timeout" if status == 504 else "error")
+        self.record(rec)
+        return _error_response(status, message, err_type, headers)
+
+
+def _error_response(status: int, message: str, err_type: str, headers: dict | None = None) -> JSONResponse:
+    return JSONResponse({"error": {"message": message, "type": err_type, "code": status}},
+                        status_code=status, headers=headers)
+
+
+def _finish_record(rec: RequestRecord, job: Job, times: list[float], toks: list[int]) -> dict:
+    rec.completion_tokens = len(toks)
+    if times:
+        rec.ttft_s = round(times[0] - job.arrival, 6)
+        rec.latency_s = round(times[-1] - job.arrival, 6)
+        p50, p95 = itl_stats(times)
+        rec.itl_p50_s = None if p50 is None else round(p50, 6)
+        rec.itl_p95_s = None if p95 is None else round(p95, 6)
+    gaps = [b - a for a, b in zip(times, times[1:])]
+    return {"queue_wait_s": rec.queue_wait_s, "ttft_s": rec.ttft_s,
+            "max_gap_s": round(max(gaps), 6) if gaps else None, "tokens": len(toks)}
+
+
+def create_app(variants: list[Variant], cfg: ServerConfig | None = None,
+               metrics: Metrics | None = None, shipper=None) -> FastAPI:
+    """`shipper`: a logship.LogShipper; its last upload runs after the drain."""
+    cfg = cfg or ServerConfig()
+    metrics = metrics or variants[0].worker.metrics
+    srv = Server(variants, cfg, metrics)
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        for v in variants:
+            v.worker.start()
+        srv.started = True
+        log.info("server ready", extra={"fields": {"variants": {v.name: v.version for v in variants},
+                                                   "slo_ttft_s": metrics.slo_ttft_s, "salt": srv.salt_source,
+                                                   "proxy_hops": cfg.proxy_hops if cfg.trust_proxy else 0}})
+        yield
+        srv.started = False
+        log.info("shutting down: draining", extra={"fields": {"drain_s": cfg.drain_s}})
+        await asyncio.gather(*(asyncio.to_thread(v.worker.stop, cfg.drain_s) for v in variants))
+        log.info("stopped")
+        if shipper is not None:
+            await asyncio.to_thread(shipper.close, cfg.drain_s)
+
+    app = FastAPI(title="tinyllm", lifespan=lifespan)
+    app.state.srv = srv
+    app.add_middleware(HeadersMiddleware, default_version=srv.router.primary.version)
+
+    @app.exception_handler(RequestValidationError)
+    async def bad_request(request: HTTPRequest, exc: RequestValidationError):
+        msg = "; ".join(f"{'.'.join(str(p) for p in e['loc'][1:])}: {e['msg']}" for e in exc.errors())
+        rec = RequestRecord(request.state.request_id, 400, client=srv.client_id(request),
+                            arrival_unix=time.time())
+        return srv.error(rec, 400, msg or "invalid request")
+
+    @app.get("/health")
+    async def health():
+        dead = [v.name for v in variants if v.worker.failed or (srv.started and not v.worker.alive)]
+        if dead:
+            return JSONResponse({"status": "failed", "variants": dead}, status_code=503)
+        return {"status": "ok"}
+
+    @app.get("/ready")
+    async def ready():
+        why = []
+        if not srv.started:
+            why.append("not started")
+        for v in variants:
+            if v.worker.failed:
+                why.append(f"{v.name}: {v.worker.failed}")
+            elif not v.worker.accepting:
+                why.append(f"{v.name}: draining")
+        if why:
+            return JSONResponse({"ready": False, "reasons": why}, status_code=503)
+        return {"ready": True}
+
+    @app.get("/v1/models")
+    async def models():
+        return {"object": "list", "data": [{"id": v.name, "object": "model", "owned_by": "tinyllm",
+                                            "version": v.version, "max_len": v.max_len} for v in variants]}
+
+    @app.get("/stats")
+    async def stats():
+        out = metrics.stats()
+        out["variants"] = {v.name: {"version": v.version, "outstanding": v.worker.outstanding(),
+                                    "n_slots": v.worker.engine.n_slots, "failed": v.worker.failed,
+                                    "norm_backend": norm_backend(v.worker.engine.model)}
+                           for v in variants}
+        return out
+
+    history_cache: dict = {}
+
+    @app.get("/stats/history")
+    async def stats_history():
+        """For the status page: computed at most once a second, however many pages poll it."""
+        now = time.monotonic()
+        if history_cache.get("t", -1.0) < now - 1.0:
+            out = metrics.history()
+            out["variants"] = {v.name: {"version": v.version, "n_slots": v.worker.engine.n_slots,
+                                        "norm_backend": norm_backend(v.worker.engine.model)}
+                               for v in variants}
+            history_cache.update(t=now, out=out)
+        return history_cache["out"]
+
+    @app.get("/", response_class=HTMLResponse)
+    async def demo_page():
+        return HTMLResponse(_page("demo.html"))
+
+    @app.get("/status", response_class=HTMLResponse)
+    async def status_page():
+        return HTMLResponse(_page("status.html"))
+
+    @app.get("/metrics")
+    async def prometheus():
+        return PlainTextResponse(metrics.prometheus(), media_type="text/plain; version=0.0.4")
+
+    @app.get("/debug/state")
+    async def debug_state(request: HTTPRequest):
+        if not request.client or request.client.host not in ("127.0.0.1", "::1"):
+            return _error_response(404, "not found", "not_found")
+        return {v.name: await v.worker.debug_state() for v in variants}
+
+    @app.post("/v1/completions")
+    async def completions(body: CompletionRequest, request: HTTPRequest):
+        arrival = time.monotonic()
+        client = srv.client_id(request)
+        rec = RequestRecord(request.state.request_id, 0, client=client, arrival_unix=time.time(),
+                            stream=body.stream)
+        variant = srv.router.pick(client)
+        rec.variant, rec.model_version = variant.name, variant.version
+        request.state.model_version = variant.version
+        worker = variant.worker
+
+        wait = srv.limiter.check(client, arrival)
+        if wait:
+            return srv.error(rec, 429, "rate limit exceeded", {"Retry-After": str(math.ceil(wait))},
+                             "rate_limit_error")
+        ids = body.prompt if isinstance(body.prompt, list) else variant.tokenizer.encode(body.prompt)
+        rec.prompt_tokens = len(ids)
+        if not ids:
+            return srv.error(rec, 400, "prompt is empty")
+        if any(not 0 <= t < variant.vocab_size for t in ids):
+            return srv.error(rec, 400, f"token ids must be in [0, {variant.vocab_size})")
+        if len(ids) + body.max_tokens > variant.max_len:
+            return srv.error(rec, 400, f"prompt ({len(ids)} tokens) + max_tokens ({body.max_tokens}) "
+                                       f"exceeds the context length {variant.max_len}")
+        rec.queue_depth_at_arrival, rec.busy_slots_at_arrival = worker.load()
+        sampling = Sampling(body.temperature, body.top_k, body.top_p)
+        eos = None if body.ignore_eos else variant.tokenizer.eos_id
+        job = Job(rec.request_id, Request(ids, body.max_tokens, eos, sampling), asyncio.get_running_loop(),
+                  asyncio.Queue(), arrival, arrival + cfg.timeout_s)
+        refused = worker.submit(job)
+        if refused == "full":
+            return srv.error(rec, 429, "server busy: queue full, retry later", {"Retry-After": "1"},
+                             "server_busy")
+        if refused:
+            return srv.error(rec, 503, "engine failed" if refused == "failed" else "server shutting down",
+                             err_type="server_error")
+
+        created, cid = int(time.time()), f"cmpl-{rec.request_id}"
+
+        async def items(poll: HTTPRequest | None):
+            """The job's items. With `poll`, also check every 0.25 s whether the client left
+            (busy requests too: tokens arriving must not hide a disconnect)."""
+            checked = time.monotonic()
+            while True:
+                if poll is None:
+                    item = await job.out.get()
+                else:
+                    try:
+                        item = await asyncio.wait_for(job.out.get(), 0.25)
+                    except TimeoutError:
+                        item = None
+                    if time.monotonic() - checked >= 0.25:
+                        checked = time.monotonic()
+                        if await poll.is_disconnected():
+                            raise _ClientGone
+                    if item is None:
+                        continue
+                if item.kind == "token" and item.queue_wait is not None:
+                    rec.queue_wait_s = round(item.queue_wait, 6)
+                yield item
+                if item.kind == "error" or item.finish:
+                    return
+
+        def usage(n: int) -> dict:
+            return {"prompt_tokens": len(ids), "completion_tokens": n, "total_tokens": len(ids) + n}
+
+        if not body.stream:
+            toks, times, finish = [], [], None
+            try:
+                async for it in items(request):
+                    if it.kind == "error":
+                        timeline = _finish_record(rec, job, times, toks)
+                        return srv.error(rec, it.status, it.message, err_type="server_error"
+                                         if it.status >= 500 else "invalid_request_error")
+                    toks.append(it.token)
+                    times.append(it.t)
+                    finish = it.finish
+            except _ClientGone:
+                worker.abort(job)
+                _finish_record(rec, job, times, toks)
+                rec.status, rec.finish_reason = 499, "abort"
+                srv.record(rec)
+                return PlainTextResponse("", status_code=499)
+            timeline = _finish_record(rec, job, times, toks)
+            rec.status, rec.finish_reason = 200, finish
+            srv.record(rec, timeline)
+            choice = {"index": 0, "text": decode_completion(variant.tokenizer, toks, ids[-1:]), "logprobs": None,
+                      "finish_reason": finish}
+            if body.return_token_ids:
+                choice["token_ids"] = toks
+            return {"id": cid, "object": "text_completion", "created": created, "model": variant.name,
+                    "choices": [choice], "usage": usage(len(toks))}
+
+        def chunk(text: str, finish: str | None, token_ids: list[int] | None = None, **extra) -> str:
+            choice = {"index": 0, "text": text, "logprobs": None, "finish_reason": finish}
+            if token_ids is not None:
+                choice["token_ids"] = token_ids
+            d = {"id": cid, "object": "text_completion", "created": created, "model": variant.name,
+                 "choices": [choice], **extra}
+            return f"data: {json.dumps(d)}\n\n"
+
+        async def stream():
+            toks, times, detok, done = [], [], _Detokenizer(variant.tokenizer, ids[-1:]), False
+            try:
+                async for it in items(None):          # disconnect: Starlette cancels this generator
+                    if it.kind == "error":
+                        rec.status, rec.error = it.status, it.message
+                        rec.finish_reason = "timeout" if it.status == 504 else "error"
+                        done = True
+                        yield (f"data: {json.dumps({'error': {'message': it.message, 'type': 'server_error', 'code': it.status}})}\n\n"
+                               "data: [DONE]\n\n")
+                        return
+                    toks.append(it.token)
+                    times.append(it.t)
+                    text = detok.add(it.token)
+                    ids_ = [it.token] if body.return_token_ids else None
+                    if it.finish:
+                        rec.status, rec.finish_reason, done = 200, it.finish, True
+                        yield chunk(text + detok.flush(), it.finish, ids_, usage=usage(len(toks)))
+                        yield "data: [DONE]\n\n"
+                        return
+                    if text or ids_:
+                        yield chunk(text, None, ids_)
+            finally:
+                if not done:
+                    worker.abort(job)
+                    rec.status, rec.finish_reason = 499, "abort"
+                timeline = _finish_record(rec, job, times, toks)
+                srv.record(rec, timeline)
+
+        return StreamingResponse(stream(), media_type="text/event-stream",
+                                 headers={"Cache-Control": "no-cache"})
+
+    return app
+
+
+# -- loading and CLI --------------------------------------------------------------------
+
+class StoryTokenizer:
+    """Our 30M model's TinyStories BPE (no BOS; EOS separates stories)."""
+
+    def __init__(self, path: Path | None = None):
+        from tinyllm.tokenizer import EOS, TOKENIZER_PATH, load_tokenizer
+        self.tok = load_tokenizer(path or TOKENIZER_PATH)
+        self.eos_id = self.tok.token_to_id(EOS)
+
+    def encode(self, text: str) -> list[int]:
+        return self.tok.encode(text).ids
+
+    def decode(self, ids: list[int]) -> str:
+        return self.tok.decode(ids)
+
+
+def git_commit() -> str:
+    if c := os.environ.get("TINYLLM_GIT_COMMIT"):
+        return c[:7]
+    try:
+        out = subprocess.run(["git", "rev-parse", "--short", "HEAD"], capture_output=True, text=True,
+                             timeout=5, cwd=Path(__file__).resolve().parent)
+        return out.stdout.strip() or "nogit"
+    except (OSError, subprocess.SubprocessError):
+        return "nogit"
+
+
+def _file_hash(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for b in iter(lambda: f.read(1 << 20), b""):
+            h.update(b)
+    return h.hexdigest()[:8]
+
+
+HF_MODEL_FILES = ["config.json", "generation_config.json", "model.safetensors", "tokenizer.json",
+                  "tokenizer_config.json"]
+
+
+def resolve_hf_model(ref: str, revision: str | None = None, local_only: bool = False) -> Path:
+    """--hf-model: a local folder as is, else a Hub repo id downloaded to the HF cache (6.5b)."""
+    if Path(ref).is_dir():
+        if revision:
+            raise ValueError("--revision is for a Hub repo id, not a local folder")
+        return Path(ref)
+    from huggingface_hub import snapshot_download
+    return Path(snapshot_download(ref, revision=revision, allow_patterns=HF_MODEL_FILES,
+                                  local_files_only=local_only))
+
+
+def load_story_hf(path: Path, device: str) -> tuple[torch.nn.Module, Tokenizer]:
+    """Our story model as exported by scripts/export_hf.py (4.8b): fp32, as trained (like --ckpt).
+
+    Only that export: its tokenizer is our TinyStories BPE (no BOS, EOS <|endoftext|>), which
+    StoryTokenizer reads from the folder. Any other Llama folder would be tokenized wrongly.
+    """
+    from tinyllm.hf_llama import load_from_dir
+    from tinyllm.tokenizer import EOS
+    tc = json.loads((path / "tokenizer_config.json").read_text())
+    if tc.get("eos_token") != EOS or tc.get("add_bos_token") or tc.get("bos_token"):
+        raise ValueError(f"{path}: not a tinyllm story-model export (eos {tc.get('eos_token')!r}, "
+                         f"bos {tc.get('bos_token')!r}); for TinyLlama use --model tinyllama")
+    tok = StoryTokenizer(path / "tokenizer.json")
+    eos_cfg = json.loads((path / "config.json").read_text()).get("eos_token_id")
+    if tok.eos_id != eos_cfg:
+        raise ValueError(f"{path}: tokenizer EOS id {tok.eos_id} != config.json eos_token_id {eos_cfg}")
+    return load_from_dir(path, device=device, dtype=torch.float32), tok
+
+
+def load_variant(args) -> tuple[torch.nn.Module, Tokenizer, str, str]:
+    """(model, tokenizer, name, version) from --model tinyllama, --ckpt or --hf-model."""
+    from tinyllm.model import Transformer
+    device = args.device
+    if args.hf_model:
+        path = resolve_hf_model(args.hf_model, args.revision, args.offline)
+        model, tok = load_story_hf(path, device)
+        name = args.hf_model.rstrip("/").split("/")[-1]
+        return model, tok, name, f"{name}@{_file_hash(path / 'model.safetensors')}+{git_commit()}"
+    if args.ckpt:
+        from tinyllm.config import ModelConfig
+        path = Path(args.ckpt)
+        ckpt = torch.load(path, map_location=device, weights_only=True)
+        model = Transformer(ModelConfig.from_dict(ckpt["model_cfg"])).to(device)
+        model.load_state_dict(ckpt["model"])
+        name = f"tinyllm-{path.parent.name}"
+        return model.eval(), StoryTokenizer(), name, f"{name}@{_file_hash(path)}+{git_commit()}"
+    from tinyllm.hf_llama import TINYLLAMA_REV, HFTokenizer, download, load_from_dir
+    path = download(local_only=args.offline)
+    model = load_from_dir(path, device=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32)
+    return model, HFTokenizer(path), "tinyllama-1.1b-chat", f"tinyllama-chat-v1.0@{TINYLLAMA_REV[:7]}+{git_commit()}"
+
+
+def main(argv: list[str] | None = None) -> None:
+    p = argparse.ArgumentParser(description="tinyllm OpenAI-compatible server (/v1/completions)")
+    src = p.add_mutually_exclusive_group(required=True)
+    src.add_argument("--model", choices=["tinyllama"])
+    src.add_argument("--ckpt", help="our checkpoint, e.g. checkpoints/main/best.pt")
+    src.add_argument("--hf-model", metavar="DIR_OR_REPO",
+                     help="our model in HF format (scripts/export_hf.py): a folder, or a Hub repo id")
+    p.add_argument("--revision", default=None, help="with a --hf-model repo id: branch, tag or commit")
+    p.add_argument("--offline", action="store_true", help="Hub models from the HF cache only")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--host", default="127.0.0.1")
+    p.add_argument("--port", type=int, default=8000)
+    p.add_argument("--slots", type=int, default=8)
+    p.add_argument("--max-len", type=int, default=None, help="per-slot context; default the model's")
+    p.add_argument("--kv", choices=["slots", "paged"], default="slots")
+    p.add_argument("--cuda-graph", action="store_true")
+    p.add_argument("--seed", type=int, default=None)
+    p.add_argument("--max-queue", type=int, default=32)
+    p.add_argument("--timeout", type=float, default=60.0)
+    p.add_argument("--rate-limit", type=float, default=0.0, help="requests/min per client; 0 = off")
+    p.add_argument("--rate-burst", type=int, default=10)
+    p.add_argument("--slow-s", type=float, default=10.0)
+    p.add_argument("--slo-ttft-s", type=float, default=1.0,
+                   help="TTFT SLO threshold on /status (SERVING.md: 1 s local GPU, 2 s on the Space)")
+    p.add_argument("--trust-proxy", action="store_true", help="client IP from X-Forwarded-For")
+    p.add_argument("--threads", type=int, default=None,
+                   help="CPU only: torch threads; default the container's CPU limit (cpu_limit), else torch's")
+    p.add_argument("--proxy-hops", type=int, default=1,
+                   help="with --trust-proxy: trusted proxies in front (client = that many entries from the right)")
+    p.add_argument("--drain-s", type=float, default=30.0)
+    p.add_argument("--log-level", default="info", choices=["debug", "info", "warning", "error"])
+    p.add_argument("--log-file", default=None, help="also write JSON lines here")
+    p.add_argument("--log-dir", default=None,
+                   help="write the shipped log lines (no text, no IP) as hourly JSONL here (logship.py)")
+    p.add_argument("--log-dataset", default=None, metavar="REPO_ID",
+                   help="with --log-dir: upload it to this private HF dataset every 10 min")
+    p.add_argument("--log-replica", default=None, help="name in the log file names; default $SPACE_ID or local")
+    args = p.parse_args(argv)
+    if args.log_dataset and not args.log_dir:
+        p.error("--log-dataset needs --log-dir")
+    if args.proxy_hops < 1:
+        p.error("--proxy-hops must be at least 1")
+    if args.threads is not None and args.threads < 1:
+        p.error("--threads must be at least 1")
+    if args.device == "cpu":
+        # an explicit --threads as given; else never more than torch's own choice (no SMT oversubscription)
+        limit = cpu_limit()
+        n = args.threads or (min(limit, torch.get_num_threads()) if limit else None)
+        if n:
+            torch.set_num_threads(n)
+
+    setup_logging(args.log_level, args.log_file)
+    shipper = None
+    if args.log_dir:
+        from tinyllm.logship import LogShipper
+        shipper = LogShipper(args.log_dir, args.log_dataset, args.log_replica).attach()
+        shipper.start_upload()
+    t = time.perf_counter()
+    model, tok, name, version = load_variant(args)
+    gen = torch.Generator(args.device).manual_seed(args.seed) if args.seed is not None else None
+    engine = ContinuousEngine(model, args.slots, args.max_len, generator=gen, kv=args.kv,
+                              cuda_graph=args.cuda_graph)
+    log.info("model loaded", extra={"fields": {"variant": name, "version": version, "device": args.device,
+                                               "slots": args.slots, "max_len": engine.max_len, "kv": args.kv,
+                                               "cuda_graph": args.cuda_graph, "threads": torch.get_num_threads(),
+                                               "load_s": round(time.perf_counter() - t, 2)}})
+    cfg = ServerConfig(max_queue=args.max_queue, timeout_s=args.timeout, rate_per_min=args.rate_limit,
+                       rate_burst=args.rate_burst, slow_s=args.slow_s, trust_proxy=args.trust_proxy,
+                       proxy_hops=args.proxy_hops, drain_s=args.drain_s)
+    metrics = Metrics(slo_ttft_s=args.slo_ttft_s)
+    worker = EngineWorker(engine, metrics, cfg.max_queue, cfg.snapshot_s, name=name)
+    app = create_app([Variant(name, version, worker, tok, engine.max_len, model.cfg.vocab_size)], cfg, metrics,
+                     shipper)
+
+    import uvicorn
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning", access_log=False,
+                timeout_graceful_shutdown=int(args.drain_s))
+
+
+if __name__ == "__main__":
+    main()

src/tinyllm/tokenizer.py

@@ -0,0 +1,35 @@
+"""Byte-level BPE tokenizer, vocab 8192 (D-007).
+
+Byte-level means every string can be encoded (unknown characters fall back to
+raw bytes), and decode(encode(text)) == text exactly.
+"""
+
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+
+from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
+
+VOCAB_SIZE = 8192
+EOS = "<|endoftext|>"  # separates stories in the token stream; id 0
+TOKENIZER_PATH = Path(__file__).resolve().parents[2] / "tokenizer" / "tinystories-bpe8192.json"
+
+
+def train_tokenizer(texts: Iterable[str] | Iterator[list[str]], vocab_size: int = VOCAB_SIZE) -> Tokenizer:
+    tok = Tokenizer(models.BPE())
+    # GPT-2-style split into words/punctuation/spaces before BPE, no prefix space
+    # (so the roundtrip is exact).
+    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tok.decoder = decoders.ByteLevel()
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=2,
+        special_tokens=[EOS],
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),  # all 256 bytes always present
+        show_progress=True,
+    )
+    tok.train_from_iterator(texts, trainer=trainer)
+    return tok
+
+
+def load_tokenizer(path: Path = TOKENIZER_PATH) -> Tokenizer:
+    return Tokenizer.from_file(str(path))

src/tinyllm/train.py

@@ -0,0 +1,278 @@
+"""Training loop (Phase 3, hyperparameters in DECISIONS.md D-019).
+
+- Step 3.2: AdamW, warmup + cosine LR, gradient accumulation, clipping, bf16 autocast.
+- Step 3.3: val loss every eval_every steps, checkpoint save/resume, sample text.
+
+The logging backend (3.4) and MFU (3.5) are added on top.
+
+    uv run python -m tinyllm.train --preset smoke --max_steps 2000
+    uv run python -m tinyllm.train --preset main --stop_at 3000      # train in chunks...
+    uv run python -m tinyllm.train --preset main --resume            # ...and continue later
+"""
+
+import argparse
+import math
+import os
+import time
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+from tinyllm.config import ModelConfig, get_config
+from tinyllm.data import DATA_DIR, get_batch, load_tokens
+from tinyllm.generate import sample
+from tinyllm.logger import Logger, make_logger
+from tinyllm.model import Transformer
+from tinyllm.tokenizer import EOS, load_tokenizer
+
+REPO = DATA_DIR.parent
+EVAL_SEED = 1234                       # same val windows at every eval, so losses are comparable
+SAMPLE_PROMPT = "Once upon a time"
+
+
+@dataclass
+class TrainConfig:
+    preset: str = "smoke"
+    seq_len: int = 512
+    micro_batch: int = 32            # sequences per forward/backward
+    grad_accum: int = 8              # micro-batches per optimizer step: 32 x 512 x 8 = 131,072 tokens/step
+    max_steps: int = 2000
+    lr: float = 1e-3                 # peak LR, reached at the end of warmup
+    min_lr: float = 1e-4             # cosine floor (10% of peak)
+    warmup_steps: int = 200
+    weight_decay: float = 0.1        # matrices and embeddings only, not norm weights
+    beta1: float = 0.9
+    beta2: float = 0.95
+    grad_clip: float = 1.0
+    seed: int = 0
+    device: str = "cuda"
+    log_every: int = 10
+    eval_every: int = 200            # val loss + checkpoint (+ sample) every N steps and at the end
+    eval_batches: int = 20           # 20 x 32 x 512 = 327,680 val tokens per eval
+    sample_tokens: int = 200         # tokens generated at each eval; 0 = no samples
+    out_dir: str = ""                # default: checkpoints/<preset>
+    resume: bool = False             # continue from <out_dir>/ckpt.pt
+    stop_at: int = 0                 # stop after this many steps (0 = max_steps); the LR schedule still uses max_steps
+    logger: str = "wandb"            # "wandb" or "none" (D-013)
+    wandb_project: str = "tinyllm"
+    wandb_mode: str = "online"       # "offline" writes to ./wandb; upload later with `wandb sync`
+    run_name: str = ""               # default: <preset>
+    peak_tflops: float = 42.0        # MFU denominator: theoretical bf16 dense peak (docs/HARDWARE.md)
+    compile: bool = True             # torch.compile the training step, +57% tok/s on main (D-020); eval + samples stay eager
+
+    @property
+    def tokens_per_step(self) -> int:
+        return self.micro_batch * self.grad_accum * self.seq_len
+
+    @property
+    def run_dir(self) -> Path:
+        return Path(self.out_dir) if self.out_dir else REPO / "checkpoints" / self.preset
+
+
+def lr_at(step: int, cfg: TrainConfig) -> float:
+    """Linear warmup to cfg.lr, then cosine decay to cfg.min_lr at cfg.max_steps."""
+    if step < cfg.warmup_steps:
+        return cfg.lr * (step + 1) / cfg.warmup_steps
+    if step >= cfg.max_steps:
+        return cfg.min_lr
+    progress = (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps)
+    return cfg.min_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (cfg.lr - cfg.min_lr)
+
+
+def make_optimizer(model: nn.Module, cfg: TrainConfig) -> torch.optim.AdamW:
+    """AdamW with weight decay on 2-D weights only. Norm gains (1-D) are not decayed."""
+    decay = [p for p in model.parameters() if p.dim() >= 2]
+    no_decay = [p for p in model.parameters() if p.dim() < 2]
+    groups = [
+        {"params": decay, "weight_decay": cfg.weight_decay},
+        {"params": no_decay, "weight_decay": 0.0},
+    ]
+    fused = torch.device(cfg.device).type == "cuda"  # one kernel for all params instead of a loop
+    return torch.optim.AdamW(groups, lr=cfg.lr, betas=(cfg.beta1, cfg.beta2), eps=1e-8, fused=fused)
+
+
+def train_step(
+    model: nn.Module,
+    opt: torch.optim.Optimizer,
+    batches: list[tuple[Tensor, Tensor]],
+    lr: float,
+    grad_clip: float,
+) -> tuple[Tensor, Tensor]:
+    """One optimizer step over len(batches) micro-batches. Returns (mean loss, pre-clip grad norm).
+
+    Each micro-batch loss is divided by the number of micro-batches, so the summed
+    gradient equals the gradient of one big batch. Results stay on the device;
+    calling .item() forces a GPU sync, so the caller only does it when logging.
+    """
+    for g in opt.param_groups:
+        g["lr"] = lr
+    device_type = batches[0][0].device.type
+    total = torch.zeros((), device=batches[0][0].device)
+    for x, y in batches:
+        with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"):
+            _, loss = model(x, targets=y)
+        (loss / len(batches)).backward()
+        total += loss.detach()
+    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+    opt.step()
+    opt.zero_grad(set_to_none=True)
+    return total / len(batches), grad_norm
+
+
+@torch.no_grad()
+def evaluate(model: nn.Module, tokens: np.ndarray, cfg: TrainConfig) -> float:
+    """Mean loss over cfg.eval_batches fixed val batches (same windows every call)."""
+    model.eval()
+    rng = np.random.default_rng(EVAL_SEED)
+    device_type = torch.device(cfg.device).type
+    total = torch.zeros((), device=cfg.device)
+    for _ in range(cfg.eval_batches):
+        x, y = get_batch(tokens, cfg.micro_batch, cfg.seq_len, rng, cfg.device)
+        with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"):
+            _, loss = model(x, targets=y)
+        total += loss
+    model.train()
+    return (total / cfg.eval_batches).item()
+
+
+def save_checkpoint(path: Path, state: dict) -> None:
+    """Write to a temp file, then rename, so a crash mid-save never corrupts the last checkpoint."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(".tmp")
+    torch.save(state, tmp)
+    os.replace(tmp, path)
+
+
+def load_checkpoint(path: Path, device: str) -> dict:
+    return torch.load(path, map_location=device, weights_only=True)
+
+
+def train(
+    cfg: TrainConfig,
+    model_cfg: ModelConfig | None = None,
+    tokens: np.ndarray | None = None,
+    val_tokens: np.ndarray | None = None,
+    log=print,
+    logger: Logger | None = None,
+) -> tuple[Transformer, list[dict]]:
+    """Run optimizer steps up to cfg.stop_at (or cfg.max_steps). Returns the model and the history.
+
+    History rows are {"step", "loss", "lr", "grad_norm", "tok_per_s"} at log steps and
+    {"step", "val_loss"} at eval steps. `log` prints to the console; `logger` (built
+    from cfg.logger when not given) sends the same numbers to W&B.
+    """
+    run_dir = cfg.run_dir
+    ckpt = load_checkpoint(run_dir / "ckpt.pt", cfg.device) if cfg.resume else None
+    if ckpt is not None:
+        model_cfg = ModelConfig.from_dict(ckpt["model_cfg"])  # the checkpoint defines the model
+    model_cfg = model_cfg or get_config(cfg.preset)
+    if cfg.seq_len > model_cfg.max_seq_len:
+        raise ValueError(f"seq_len {cfg.seq_len} > model max_seq_len {model_cfg.max_seq_len}")
+    torch.manual_seed(cfg.seed)
+    torch.set_float32_matmul_precision("high")  # TF32 for any matmul left in fp32
+    tokens = load_tokens("train") if tokens is None else tokens
+    val_tokens = load_tokens("validation") if val_tokens is None else val_tokens
+    tok = load_tokenizer() if cfg.sample_tokens else None
+    rng = np.random.default_rng(cfg.seed)
+
+    model = Transformer(model_cfg).to(cfg.device)
+    opt = make_optimizer(model, cfg)
+    start, best_val = 0, math.inf
+    if ckpt is not None:
+        model.load_state_dict(ckpt["model"])
+        opt.load_state_dict(ckpt["optimizer"])
+        rng.bit_generator.state = ckpt["rng"]
+        start, best_val = ckpt["step"] + 1, ckpt["best_val"]
+        log(f"resumed from {run_dir / 'ckpt.pt'} at step {start} (best val {best_val:.4f})")
+    end = min(cfg.stop_at or cfg.max_steps, cfg.max_steps)
+    log(f"params {model.num_params():,} | {cfg.tokens_per_step:,} tokens/step "
+        f"({cfg.micro_batch}x{cfg.seq_len} x {cfg.grad_accum} accum) | steps {start}..{end - 1} "
+        f"of {cfg.max_steps} ({cfg.tokens_per_step * cfg.max_steps / 1e6:.1f}M tokens total) | {run_dir}")
+    if logger is None:
+        logger = make_logger(
+            cfg.logger, project=cfg.wandb_project, name=cfg.run_name or cfg.preset,
+            config={**asdict(cfg), "model": model_cfg.to_dict(), "params": model.num_params()},
+            mode=cfg.wandb_mode, run_id=ckpt.get("wandb_id") if ckpt else None, dir=str(REPO))
+    try:
+        history = _loop(cfg, model, opt, rng, tokens, val_tokens, tok, model_cfg, run_dir,
+                        start, end, best_val, log, logger)
+    finally:
+        logger.finish()  # flush W&B even if training crashes or is interrupted
+    return model, history
+
+
+def _loop(cfg, model, opt, rng, tokens, val_tokens, tok, model_cfg, run_dir,
+          start, end, best_val, log, logger) -> list[dict]:
+    history: list[dict] = []
+    flops_per_token = model_cfg.flops_per_token(cfg.seq_len)
+    # The compiled wrapper shares the model's weights. Only the fixed-shape training step uses it:
+    # eval and sampling (a new length every token) would trigger recompiles, and checkpoints
+    # save `model` so state_dict keys never get the `_orig_mod.` prefix.
+    step_model = torch.compile(model) if cfg.compile else model
+    t_last, step_last = time.perf_counter(), start - 1
+    for step in range(start, end):
+        lr = lr_at(step, cfg)
+        batches = [get_batch(tokens, cfg.micro_batch, cfg.seq_len, rng, cfg.device)
+                   for _ in range(cfg.grad_accum)]
+        loss, grad_norm = train_step(step_model, opt, batches, lr, cfg.grad_clip)
+
+        if step % cfg.log_every == 0 or step == end - 1:
+            rec = {"step": step, "loss": loss.item(), "lr": lr, "grad_norm": grad_norm.item()}
+            now = time.perf_counter()  # .item() above synced the GPU, so this timing is real
+            rec["tok_per_s"] = (step - step_last) * cfg.tokens_per_step / (now - t_last)
+            rec["tflops"] = rec["tok_per_s"] * flops_per_token / 1e12   # achieved model TFLOPS
+            rec["mfu"] = rec["tflops"] / cfg.peak_tflops
+            t_last, step_last = now, step
+            history.append(rec)
+            log(f"step {step:5d} | loss {rec['loss']:.4f} | lr {lr:.2e} | "
+                f"grad {rec['grad_norm']:.3f} | {rec['tok_per_s'] / 1e3:.1f}K tok/s | "
+                f"{rec['tflops']:.1f} TFLOPS | MFU {rec['mfu']:.1%}")
+            logger.log({"train/loss": rec["loss"], "train/lr": lr, "train/grad_norm": rec["grad_norm"],
+                        "train/tok_per_s": rec["tok_per_s"], "train/tflops": rec["tflops"],
+                        "train/mfu": rec["mfu"],
+                        "train/tokens": (step + 1) * cfg.tokens_per_step}, step=step)
+
+        if step % cfg.eval_every == 0 or step == end - 1:
+            val = evaluate(model, val_tokens, cfg)
+            history.append({"step": step, "val_loss": val})
+            logger.log({"val/loss": val, "val/best": min(val, best_val)}, step=step)
+            saved = "ckpt.pt"
+            if val < best_val:
+                best_val = val
+                # Weights only, for sampling / export later.
+                save_checkpoint(run_dir / "best.pt", {"model": model.state_dict(),
+                                "model_cfg": model_cfg.to_dict(), "step": step, "val_loss": val})
+                saved += " + best.pt"
+            save_checkpoint(run_dir / "ckpt.pt", {
+                "model": model.state_dict(), "optimizer": opt.state_dict(), "rng": rng.bit_generator.state,
+                "step": step, "best_val": best_val, "model_cfg": model_cfg.to_dict(), "train_cfg": asdict(cfg),
+                "wandb_id": logger.run_id,
+            })
+            log(f"eval  {step:5d} | val {val:.4f} | best {best_val:.4f} | saved {saved}")
+            if tok is not None:
+                gen = torch.Generator(cfg.device).manual_seed(cfg.seed)  # same sampling noise each eval
+                prompt = torch.tensor([tok.encode(SAMPLE_PROMPT).ids], device=cfg.device)
+                out = sample(model, prompt, cfg.sample_tokens, temperature=0.8, top_k=50,
+                             eos_id=tok.token_to_id(EOS), generator=gen)
+                text = tok.decode(out[0].tolist())
+                log("sample | " + text.replace("\n", " "))
+                logger.log_sample(text, step=step)
+            t_last, step_last = time.perf_counter(), step  # eval time doesn't count toward tok/s
+    return history
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Train a tinyllm model")
+    for f in fields(TrainConfig):
+        if isinstance(f.default, bool):
+            p.add_argument(f"--{f.name}", action=argparse.BooleanOptionalAction, default=f.default)
+        else:
+            p.add_argument(f"--{f.name}", type=type(f.default), default=f.default)
+    train(TrainConfig(**vars(p.parse_args())))
+
+
+if __name__ == "__main__":
+    main()

src/tinyllm/web/demo.html

@@ -0,0 +1,250 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>tiny-llm demo</title>
+<style>
+:root {
+  color-scheme: light;
+  --page: #f9f9f7; --surface: #fcfcfb; --ink: #0b0b0b; --ink-2: #52514e; --muted: #898781;
+  --grid: #e1e0d9; --ring: rgba(11,11,11,0.10); --accent: #2a78d6; --on-accent: #fff;
+  --gen: #e8f1fc; --bad: #d03b3b; --warn: #fab219;
+}
+@media (prefers-color-scheme: dark) {
+  :root:not([data-theme="light"]) {
+    color-scheme: dark;
+    --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781;
+    --grid: #2c2c2a; --ring: rgba(255,255,255,0.10); --accent: #3987e5; --gen: #12243a;
+  }
+}
+:root[data-theme="dark"] {
+  color-scheme: dark;
+  --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781;
+  --grid: #2c2c2a; --ring: rgba(255,255,255,0.10); --accent: #3987e5; --gen: #12243a;
+}
+* { box-sizing: border-box; }
+body { margin: 0; background: var(--page); color: var(--ink);
+       font: 16px/1.5 system-ui, -apple-system, "Segoe UI", sans-serif; }
+main { max-width: 760px; margin: 0 auto; padding: 28px 16px 48px; }
+h1 { font-size: 24px; margin: 0; font-weight: 600; }
+.lede { color: var(--ink-2); margin: 4px 0 20px; }
+a { color: var(--accent); }
+.panel { background: var(--surface); border: 1px solid var(--ring); border-radius: 12px; padding: 16px; }
+label { font-size: 14px; color: var(--ink-2); }
+textarea { width: 100%; min-height: 84px; resize: vertical; margin: 6px 0 4px; padding: 10px 12px;
+           font: inherit; color: var(--ink); background: var(--page); border: 1px solid var(--ring); border-radius: 8px; }
+textarea:focus-visible, button:focus-visible, input:focus-visible { outline: 2px solid var(--accent); outline-offset: 2px; }
+.count { font-size: 12px; color: var(--muted); text-align: right; }
+.chips { display: flex; flex-wrap: wrap; gap: 6px; margin: 6px 0 14px; }
+.chip { font: inherit; font-size: 13px; padding: 4px 10px; border-radius: 999px; cursor: pointer;
+        color: var(--ink-2); background: transparent; border: 1px solid var(--ring); }
+.chip:hover { color: var(--ink); border-color: var(--muted); }
+.controls { display: grid; grid-template-columns: 1fr 1fr; gap: 12px 20px; }
+.controls div { display: flex; flex-direction: column; }
+.controls output { color: var(--ink); font-variant-numeric: tabular-nums; }
+input[type=range] { accent-color: var(--accent); width: 100%; }
+.actions { display: flex; gap: 10px; align-items: center; margin-top: 16px; flex-wrap: wrap; }
+button.primary, button.secondary { font: inherit; font-size: 15px; padding: 8px 18px; border-radius: 8px; cursor: pointer; }
+button.primary { background: var(--accent); color: var(--on-accent); border: 1px solid var(--accent); }
+button.secondary { background: transparent; color: var(--ink); border: 1px solid var(--ring); }
+button:disabled { opacity: 0.5; cursor: default; }
+.hint { font-size: 13px; color: var(--muted); }
+.story { margin-top: 16px; min-height: 120px; white-space: pre-wrap; font-size: 17px; line-height: 1.65; }
+.story .p { color: var(--ink-2); }
+.story .g { background: var(--gen); border-radius: 3px; }
+.story .placeholder { color: var(--muted); }
+.cursor { display: inline-block; width: 8px; height: 1.1em; vertical-align: text-bottom; background: var(--accent);
+          margin-left: 1px; animation: blink 1s steps(1) infinite; }
+@keyframes blink { 50% { opacity: 0; } }
+@media (prefers-reduced-motion: reduce) { .cursor { animation: none; } }
+.stats { display: flex; flex-wrap: wrap; gap: 6px 20px; margin-top: 12px; font-size: 13px; color: var(--ink-2);
+         font-variant-numeric: tabular-nums; }
+.stats b { color: var(--ink); font-weight: 600; }
+.msg { display: none; margin-top: 12px; padding: 8px 12px; border-radius: 8px; font-size: 14px;
+       border: 1px solid var(--ring); border-left: 4px solid var(--warn); }
+.msg.show { display: block; }
+.msg.error { border-left-color: var(--bad); }
+footer { color: var(--ink-2); font-size: 13px; margin-top: 28px; }
+footer p { margin: 6px 0; }
+footer details { margin: 6px 0; }
+footer summary { cursor: pointer; color: var(--ink); }
+footer ul { margin: 6px 0; padding-left: 20px; }
+@media (max-width: 520px) { .controls { grid-template-columns: 1fr; } }
+</style>
+</head>
+<body>
+<main>
+  <h1>tiny-llm</h1>
+  <p class="lede">A small language model trained from scratch on children's stories.
+    Start a story and it writes the rest.
+    It knows simple stories, not facts, so it can't answer questions or follow instructions.</p>
+
+  <section class="panel">
+    <label for="prompt">Story opening</label>
+    <textarea id="prompt" maxlength="600" spellcheck="true">Once upon a time, there was a little dog named Max.</textarea>
+    <div class="count" id="count"></div>
+    <div class="chips" id="chips" aria-label="Example openings"></div>
+    <div class="controls">
+      <div><label for="len">Length: <output id="len-out"></output> tokens</label>
+        <input type="range" id="len" min="20" max="250" step="10" value="120"></div>
+      <div><label for="temp">Temperature: <output id="temp-out"></output></label>
+        <input type="range" id="temp" min="0" max="1.5" step="0.1" value="0.8"></div>
+    </div>
+    <div class="actions">
+      <button class="primary" id="go">Write the story</button>
+      <button class="secondary" id="stop" disabled>Stop</button>
+      <span class="hint">Ctrl + Enter. Lower temperature: safer and more repetitive; higher: more surprising.</span>
+    </div>
+  </section>
+
+  <div class="msg" id="msg" role="status"></div>
+  <div class="story" id="story" aria-live="polite"><span class="placeholder">The story appears here, one token at a time.</span></div>
+  <div class="stats" id="stats"></div>
+
+  <footer>
+    <p>Model: <span id="model">loading…</span> ·
+      <a href="https://huggingface.co/CoolDip00/tiny-llm-tinystories-30m">model card</a> ·
+      <a href="https://github.com/cooldip00/tiny-llm">source on GitHub</a> ·
+      <a href="status">server status</a></p>
+    <p>This demo runs on a shared CPU, so it is slower than the GPU numbers in the project README,
+      and it can be busy or asleep for a minute.</p>
+    <p>Privacy: your text isn't stored. The server keeps request metadata (timings, token counts and a
+      salted hash of your address), never the opening or the story.</p>
+    <details>
+      <summary>What is logged</summary>
+      <ul>
+        <li>Per request: time, status, timings (queue wait, time to first token, total), token counts,
+          whether it streamed, why it stopped, the error message if any, and how busy the server was.</li>
+        <li>Your address only as a salted hash: it can tell repeat visits apart, not who you are.
+          Raw addresses are never stored.</li>
+        <li>Never the opening you type or the story the model writes.</li>
+        <li>The log is kept in a private dataset that only the project owner can read, and is deleted after 30 days.</li>
+      </ul>
+    </details>
+  </footer>
+</main>
+<script>
+"use strict";
+const EXAMPLES = [
+  "Once upon a time, there was a little dog named Max.",
+  "Lily found a shiny key under her bed.",
+  "The sun was hot, and Tom wanted to swim.",
+  "One day, a small bird could not fly.",
+  "Mia and her mom baked a big cake.",
+];
+const MAX_CHARS = 600;
+const $ = (id) => document.getElementById(id);
+const prompt = $("prompt"), story = $("story"), stats = $("stats"), msg = $("msg");
+let ctrl = null;
+
+for (const ex of EXAMPLES) {
+  const b = document.createElement("button");
+  b.className = "chip"; b.type = "button"; b.textContent = ex;
+  b.onclick = () => { prompt.value = ex; updateCount(); prompt.focus(); };
+  $("chips").appendChild(b);
+}
+function updateCount() { $("count").textContent = `${prompt.value.length} / ${MAX_CHARS}`; }
+function bindRange(id, fmt) {
+  const r = $(id), o = $(id + "-out");
+  const f = () => { o.textContent = fmt(r.value); };
+  r.oninput = f; f();
+}
+bindRange("len", (v) => v);
+bindRange("temp", (v) => Number(v).toFixed(1));
+prompt.oninput = updateCount; updateCount();
+prompt.onkeydown = (e) => { if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) { e.preventDefault(); start(); } };
+$("go").onclick = start;
+$("stop").onclick = () => ctrl && ctrl.abort();
+
+fetch("v1/models").then((r) => r.json()).then((d) => {
+  const m = d.data && d.data[0];
+  $("model").textContent = m ? (m.version || m.id) : "unknown";
+}).catch(() => { $("model").textContent = "unavailable"; });
+
+function say(text, kind) {
+  msg.textContent = text;
+  msg.className = "msg show" + (kind === "error" ? " error" : "");
+}
+function stat(label, value) {
+  const s = document.createElement("span");
+  const b = document.createElement("b"); b.textContent = value;
+  s.append(document.createTextNode(label + " "), b);
+  return s;
+}
+function showStats(t) {
+  const items = [];
+  if (t.ttft !== null) items.push(stat("First token", Math.round(t.ttft * 1000) + " ms"));
+  if (t.tokens > 1 && t.genS > 0) items.push(stat("Speed", ((t.tokens - 1) / t.genS).toFixed(1) + " tokens/s"));
+  items.push(stat("Tokens", String(t.tokens)));
+  if (t.finish) items.push(stat("Stopped because", t.finish === "stop" ? "the story ended" :
+                                t.finish === "length" ? "it reached the length" : t.finish));
+  stats.replaceChildren(...items);
+}
+function errorText(status, body, retry) {
+  if (status === 429) return `The server is busy right now. Please try again in ${retry || "a few"} seconds.`;
+  if (status === 503) return "The model is still loading (or the server is waking up). Try again in a minute.";
+  if (status === 504) return "That took too long and was stopped. Try a shorter length.";
+  if (status === 400) return "The request wasn't accepted: " + ((body && body.error && body.error.message) || "bad input");
+  return `Something went wrong on the server (${status}). Please try again.`;
+}
+
+async function start() {
+  if (ctrl) return;
+  const text = prompt.value.slice(0, MAX_CHARS);
+  if (!text.trim()) { say("Write the start of a story first."); return; }
+  msg.className = "msg";
+  const pSpan = document.createElement("span"), gSpan = document.createElement("span"), cur = document.createElement("span");
+  pSpan.className = "p"; gSpan.className = "g"; cur.className = "cursor";
+  pSpan.textContent = text;
+  story.replaceChildren(pSpan, gSpan, cur);
+  stats.replaceChildren();
+  ctrl = new AbortController();
+  $("go").disabled = true; $("stop").disabled = false;
+  const t = { t0: performance.now(), ttft: null, tFirst: 0, genS: 0, tokens: 0, finish: null };
+  try {
+    const r = await fetch("v1/completions", {
+      method: "POST", headers: { "content-type": "application/json" }, signal: ctrl.signal,
+      body: JSON.stringify({ prompt: text, max_tokens: Number($("len").value),
+                             temperature: Number($("temp").value), stream: true }),
+    });
+    if (!r.ok) {
+      let body = null; try { body = await r.json(); } catch (e) {}
+      throw { status: r.status, body, retry: r.headers.get("retry-after") };
+    }
+    const reader = r.body.getReader(), dec = new TextDecoder();
+    let buf = "";
+    for (;;) {
+      const { value, done } = await reader.read();
+      if (done) break;
+      buf += dec.decode(value, { stream: true });
+      let i;
+      while ((i = buf.indexOf("\n\n")) >= 0) {
+        const line = buf.slice(0, i).trim(); buf = buf.slice(i + 2);
+        if (!line.startsWith("data:")) continue;
+        const data = line.slice(5).trim();
+        if (data === "[DONE]") continue;
+        const ev = JSON.parse(data);
+        if (ev.error) throw { status: ev.error.code, body: ev };
+        const ch = ev.choices[0], now = performance.now();
+        if (t.ttft === null) { t.ttft = (now - t.t0) / 1000; t.tFirst = now; }
+        t.genS = (now - t.tFirst) / 1000;
+        if (ev.usage) t.tokens = ev.usage.completion_tokens; else t.tokens += 1;
+        if (ch.text) gSpan.textContent += ch.text;
+        if (ch.finish_reason) t.finish = ch.finish_reason;
+        showStats(t);
+      }
+    }
+  } catch (e) {
+    if (e && e.name === "AbortError") { t.finish = "you pressed Stop"; showStats(t); }
+    else if (e && e.status) say(errorText(e.status, e.body, e.retry), e.status === 429 || e.status === 503 ? "warn" : "error");
+    else say("Can't reach the server. Check your connection, or try again in a minute.", "error");
+  } finally {
+    cur.remove();
+    ctrl = null;
+    $("go").disabled = false; $("stop").disabled = true;
+  }
+}
+</script>
+</body>
+</html>

src/tinyllm/web/status.html

@@ -0,0 +1,401 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>tiny-llm status</title>
+<style>
+:root {
+  color-scheme: light;
+  --page: #f9f9f7; --surface: #fcfcfb; --ink: #0b0b0b; --ink-2: #52514e; --muted: #898781;
+  --grid: #e1e0d9; --axis: #c3c2b7; --ring: rgba(11,11,11,0.10);
+  --s1: #2a78d6; --s2: #eb6834; --s3: #1baf7a; --track: #cde2fb;
+  --good: #0ca30c; --warn: #fab219; --bad: #d03b3b;
+}
+@media (prefers-color-scheme: dark) {
+  :root:not([data-theme="light"]) {
+    color-scheme: dark;
+    --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781;
+    --grid: #2c2c2a; --axis: #383835; --ring: rgba(255,255,255,0.10);
+    --s1: #3987e5; --s2: #d95926; --s3: #199e70; --track: #184f95;
+  }
+}
+:root[data-theme="dark"] {
+  color-scheme: dark;
+  --page: #0d0d0d; --surface: #1a1a19; --ink: #f0efec; --ink-2: #c3c2b7; --muted: #898781;
+  --grid: #2c2c2a; --axis: #383835; --ring: rgba(255,255,255,0.10);
+  --s1: #3987e5; --s2: #d95926; --s3: #199e70; --track: #184f95;
+}
+* { box-sizing: border-box; }
+body { margin: 0; background: var(--page); color: var(--ink);
+       font: 15px/1.45 system-ui, -apple-system, "Segoe UI", sans-serif; }
+main { max-width: 980px; margin: 0 auto; padding: 24px 16px 48px; }
+header { display: flex; flex-wrap: wrap; align-items: baseline; gap: 4px 16px; }
+h1 { font-size: 22px; margin: 0; font-weight: 600; }
+h2 { font-size: 16px; margin: 0 0 2px; font-weight: 600; }
+a { color: var(--s1); }
+.meta, .sub { color: var(--ink-2); font-size: 13px; }
+.sub { margin: 0 0 10px; }
+.banner { display: none; margin: 12px 0 0; padding: 8px 12px; border-radius: 8px;
+          border: 1px solid var(--ring); background: var(--surface); }
+.banner.show { display: block; }
+.cards { display: grid; grid-template-columns: repeat(auto-fit, minmax(260px, 1fr)); gap: 12px; margin: 20px 0 8px; }
+.card, .panel { background: var(--surface); border: 1px solid var(--ring); border-radius: 12px; padding: 14px 16px; }
+.card .value { font-size: 28px; font-weight: 600; margin: 4px 0 0; }
+.card .line { color: var(--ink-2); font-size: 13px; margin: 2px 0; }
+.state { display: inline-flex; align-items: center; gap: 6px; font-size: 13px; font-weight: 600; }
+.state i { display: inline-grid; place-items: center; width: 16px; height: 16px; border-radius: 50%;
+           font-style: normal; font-size: 11px; color: #fff; flex: none; }
+.state.none { font-weight: 400; align-items: flex-start; } .state.none i { margin-top: 2px; }
+.state.good i { background: var(--good); } .state.bad i { background: var(--bad); }
+.state.warn i { background: var(--warn); color: #0b0b0b; } .state.none i { background: var(--muted); }
+.meter { height: 8px; border-radius: 4px; background: var(--track); margin: 8px 0 4px; overflow: hidden; }
+.meter > div { height: 100%; border-radius: 4px; background: var(--s1); }
+.window { color: var(--ink-2); font-size: 13px; margin: 0 0 16px; }
+.panel { margin: 12px 0; position: relative; }
+.legend { display: flex; flex-wrap: wrap; gap: 4px 14px; font-size: 13px; color: var(--ink-2); margin: 0 0 6px; }
+.legend span { display: inline-flex; align-items: center; gap: 6px; }
+.key { display: inline-block; width: 14px; height: 0; border-top: 2px solid; }
+.key.box { width: 10px; height: 10px; border: 0; border-radius: 2px; }
+.key.ref { border-top: 1px solid var(--muted); }
+.chart { width: 100%; height: 190px; display: block; outline: none; }
+.chart:focus-visible { box-shadow: 0 0 0 2px var(--s1); border-radius: 4px; }
+.chart text { fill: var(--muted); font-size: 11px; font-variant-numeric: tabular-nums; }
+.tip { position: absolute; pointer-events: none; display: none; min-width: 140px; z-index: 2;
+       background: var(--surface); border: 1px solid var(--ring); border-radius: 8px;
+       padding: 8px 10px; font-size: 13px; box-shadow: 0 4px 16px rgba(0,0,0,0.12); }
+.tip .t { color: var(--ink-2); margin-bottom: 4px; }
+.tip .r { display: flex; align-items: center; gap: 8px; }
+.tip .r b { font-weight: 600; min-width: 52px; }
+.tip .r span { color: var(--ink-2); }
+details { margin: 16px 0; }
+summary { cursor: pointer; color: var(--ink-2); }
+table { border-collapse: collapse; width: 100%; font-size: 13px; margin-top: 8px; font-variant-numeric: tabular-nums; }
+th, td { text-align: right; padding: 4px 8px; border-bottom: 1px solid var(--grid); }
+th:first-child, td:first-child { text-align: left; }
+.tablewrap { overflow-x: auto; }
+footer { color: var(--ink-2); font-size: 13px; margin-top: 28px; }
+.stale { opacity: 0.55; }
+</style>
+</head>
+<body>
+<main>
+  <header>
+    <h1>Status</h1>
+    <span class="meta" id="meta">Loading…</span>
+  </header>
+  <p class="meta"><a href="./">← Demo</a> · numbers are measured on this server, from arrival to token sent</p>
+  <div class="banner" id="banner" role="status"></div>
+
+  <section class="cards" id="cards" aria-label="Service level objectives"></section>
+  <p class="window" id="window"></p>
+
+  <section class="panel">
+    <h2>Requests per minute</h2>
+    <p class="sub">Last 60 minutes, by outcome</p>
+    <div class="legend">
+      <span><i class="key box" style="background:var(--s1)"></i>Served</span>
+      <span><i class="key box" style="background:var(--s2)"></i>Busy (429)</span>
+      <span><i class="key box" style="background:var(--s3)"></i>Server error (5xx)</span>
+    </div>
+    <svg class="chart" id="c-req" tabindex="0" role="img" aria-label="Requests per minute, last hour"></svg>
+    <div class="tip"></div>
+  </section>
+
+  <section class="panel">
+    <h2>Latency</h2>
+    <p class="sub">Per minute, successful requests; time to first token includes the wait for a slot</p>
+    <div class="legend">
+      <span><i class="key" style="border-color:var(--s1)"></i>First token p50</span>
+      <span><i class="key" style="border-color:var(--s2)"></i>First token p95</span>
+      <span><i class="key" style="border-color:var(--s3)"></i>Whole request p95</span>
+      <span><i class="key ref"></i><span id="ref-label">SLO</span></span>
+    </div>
+    <svg class="chart" id="c-lat" tabindex="0" role="img" aria-label="Latency per minute, last hour"></svg>
+    <div class="tip"></div>
+  </section>
+
+  <section class="panel">
+    <h2>Load</h2>
+    <p class="sub">Engine snapshots every 5 s</p>
+    <div class="legend">
+      <span><i class="key" style="border-color:var(--s1)"></i>Busy slots</span>
+      <span><i class="key" style="border-color:var(--s2)"></i>Waiting in queue</span>
+      <span><i class="key ref"></i><span id="slots-label">Slots</span></span>
+    </div>
+    <svg class="chart" id="c-load" tabindex="0" role="img" aria-label="Busy slots and queue length, last hour"></svg>
+    <div class="tip"></div>
+  </section>
+
+  <details>
+    <summary>Show the last hour as a table</summary>
+    <div class="tablewrap"><table id="table"></table></div>
+  </details>
+
+  <footer>
+    SLOs and the error budget are defined in
+    <a href="https://github.com/cooldip00/tiny-llm/blob/main/docs/SERVING.md#slos-413-step-3-written-before-the-load-test">docs/SERVING.md</a>.
+    History is kept in memory, so it starts again when the server restarts.
+    The server runs on CPU here; the speed numbers in the project README come from a laptop GPU.
+  </footer>
+</main>
+<script>
+"use strict";
+const NS = "http://www.w3.org/2000/svg";
+const SLO_NAMES = { ttft: "Time to first token", availability: "Availability", shed: "Requests not shed" };
+let last = null;
+
+function el(tag, attrs, parent, text) {
+  const e = document.createElementNS(NS, tag);
+  for (const k in attrs) e.setAttribute(k, attrs[k]);
+  if (text !== undefined) e.textContent = text;
+  if (parent) parent.appendChild(e);
+  return e;
+}
+function h(tag, cls, text) {
+  const e = document.createElement(tag);
+  if (cls) e.className = cls;
+  if (text !== undefined) e.textContent = text;
+  return e;
+}
+function niceStep(v) {
+  if (!(v > 0)) return 1;
+  const p = Math.pow(10, Math.floor(Math.log10(v)));
+  for (const m of [1, 2, 5]) if (m * p >= v) return m * p;
+  return 10 * p;
+}
+const pct = (x, d = 1) => x === null || x === undefined ? "–" : (100 * x).toFixed(d) + "%";
+function secs(x) {
+  if (x === null || x === undefined) return "–";
+  return x < 1 ? Math.round(x * 1000) + " ms" : x.toFixed(2) + " s";
+}
+function dur(s) {
+  s = Math.round(s);
+  const d = Math.floor(s / 86400), hh = Math.floor(s % 86400 / 3600), m = Math.floor(s % 3600 / 60);
+  if (d) return `${d} d ${hh} h`;
+  if (hh) return `${hh} h ${m} min`;
+  return m ? `${m} min` : `${s} s`;
+}
+const minLabel = (ago) => ago === 0 ? "this minute" : `${ago} min ago`;
+
+/* One chart: x in minutes from -60 to 0; series of {name, color, values: [{x, y}|null...]}.
+   kind "stack": columns, one per minute; kind "line": 2px lines with an end dot. */
+function drawChart(svg, spec) {
+  const W = svg.clientWidth || 600, H = svg.clientHeight || 190;
+  const m = { l: 44, r: 14, t: 10, b: 24 };
+  const pw = W - m.l - m.r, ph = H - m.t - m.b;
+  svg.setAttribute("viewBox", `0 0 ${W} ${H}`);
+  svg.replaceChildren();
+  const X = (x) => m.l + (x + 60) / 60 * pw;
+  let ymax = spec.ymin || 0;
+  for (const p of spec.points) ymax = Math.max(ymax, spec.total(p));
+  if (spec.ref) ymax = Math.max(ymax, spec.ref.value * 1.15);
+  let step = niceStep(ymax / 4);                 // ticks on 1 / 2 / 5 × 10^k, 3–5 gridlines
+  if (spec.integer) step = Math.max(1, Math.round(step));
+  const nTicks = Math.max(1, Math.ceil(ymax / step - 1e-9));
+  ymax = nTicks * step;
+  const Y = (y) => m.t + ph - y / ymax * ph;
+  for (let i = 0; i <= nTicks; i++) {
+    const v = +(step * i).toPrecision(12), y = Y(v);
+    el("line", { x1: m.l, x2: W - m.r, y1: y, y2: y, stroke: i ? "var(--grid)" : "var(--axis)", "stroke-width": 1 }, svg);
+    el("text", { x: m.l - 6, y: y + 4, "text-anchor": "end" }, svg, spec.fmt(v));
+  }
+  for (const ago of [60, 45, 30, 15, 0])
+    el("text", { x: X(-ago), y: H - 6, "text-anchor": ago === 60 ? "start" : ago === 0 ? "end" : "middle" },
+       svg, ago ? `−${ago} min` : "now");
+  if (spec.ref) {
+    const y = Y(spec.ref.value);
+    el("line", { x1: m.l, x2: W - m.r, y1: y, y2: y, stroke: "var(--muted)", "stroke-width": 1 }, svg);
+    el("text", { x: W - m.r, y: y - 4, "text-anchor": "end" }, svg, spec.ref.label);
+  }
+  if (spec.kind === "stack") {
+    const slot = pw / 60, bw = Math.max(2, Math.min(24, slot - 2));
+    for (const p of spec.points) {
+      if (!p) continue;
+      let base = 0;
+      const parts = spec.series.map((s) => ({ s, v: s.get(p) })).filter((q) => q.v > 0);
+      parts.forEach((q, i) => {
+        const x = X(p.x) - slot / 2 + (slot - bw) / 2, y0 = Y(base), y1 = Y(base + q.v);
+        const top = i === parts.length - 1, gap = i ? 2 : 0, hgt = Math.max(0, y0 - y1 - gap);
+        const r = top ? Math.min(2, bw / 2, hgt) : 0, yb = y0 - gap;
+        el("path", { d: `M${x},${yb} V${yb - hgt + r} Q${x},${yb - hgt} ${x + r},${yb - hgt} H${x + bw - r} ` +
+                        `Q${x + bw},${yb - hgt} ${x + bw},${yb - hgt + r} V${yb} Z`, fill: q.s.color }, svg);
+        base += q.v;
+      });
+    }
+  } else {
+    for (const s of spec.series) {
+      let d = "", lastPt = null, prev = null;
+      for (const p of spec.points) {
+        const v = p ? s.get(p) : null;
+        if (v === null || v === undefined) { prev = null; continue; }
+        const gapBreak = prev && p.x - prev.x > (spec.maxGap || 1.01);
+        d += `${prev && !gapBreak ? "L" : "M"}${X(p.x).toFixed(1)},${Y(v).toFixed(1)} `;
+        prev = p; lastPt = [X(p.x), Y(v)];
+      }
+      if (d) el("path", { d, fill: "none", stroke: s.color, "stroke-width": 2, "stroke-linejoin": "round", "stroke-linecap": "round" }, svg);
+      if (lastPt) el("circle", { cx: lastPt[0], cy: lastPt[1], r: 4, fill: s.color, stroke: "var(--surface)", "stroke-width": 2 }, svg);
+    }
+  }
+  attachHover(svg, spec, X, m, H);
+}
+
+function attachHover(svg, spec, X, m, H) {
+  const tip = svg.nextElementSibling;
+  const pts = spec.points.filter(Boolean);
+  const cross = el("line", { y1: m.t, y2: H - m.b, stroke: "var(--axis)", "stroke-width": 1, visibility: "hidden" }, svg);
+  let idx = -1;
+  function show(i) {
+    if (!pts.length) return;
+    idx = Math.max(0, Math.min(pts.length - 1, i));
+    const p = pts[idx], x = X(p.x);
+    cross.setAttribute("x1", x); cross.setAttribute("x2", x); cross.setAttribute("visibility", "visible");
+    tip.replaceChildren(h("div", "t", spec.when(p)));
+    for (const s of spec.series) {
+      const r = h("div", "r"), k = h("i", spec.kind === "stack" ? "key box" : "key");
+      k.style[spec.kind === "stack" ? "background" : "borderColor"] = s.color;
+      r.append(k, h("b", "", spec.fmtTip(s.get(p))), h("span", "", s.name));
+      tip.appendChild(r);
+    }
+    tip.style.display = "block";
+    const box = svg.getBoundingClientRect(), panel = svg.parentElement.getBoundingClientRect();
+    let left = box.left - panel.left + x + 12;
+    if (left + tip.offsetWidth > panel.width - 8) left = box.left - panel.left + x - tip.offsetWidth - 12;
+    tip.style.left = Math.max(8, left) + "px";
+    tip.style.top = (box.top - panel.top + 8) + "px";
+  }
+  function hide() { cross.setAttribute("visibility", "hidden"); tip.style.display = "none"; }
+  function nearest(clientX) {
+    const box = svg.getBoundingClientRect(), x = clientX - box.left;
+    let best = -1, bd = Infinity;
+    pts.forEach((p, i) => { const d = Math.abs(X(p.x) - x); if (d < bd) { bd = d; best = i; } });
+    return best;
+  }
+  svg.onpointermove = (e) => show(nearest(e.clientX));
+  svg.onpointerleave = hide;
+  svg.onblur = hide;
+  svg.onfocus = () => show(pts.length - 1);
+  svg.onkeydown = (e) => {
+    if (e.key === "ArrowLeft") { show(idx - 1); e.preventDefault(); }
+    if (e.key === "ArrowRight") { show(idx + 1); e.preventDefault(); }
+    if (e.key === "Escape") hide();
+  };
+}
+
+function state(cls, icon, label) {
+  const s = h("span", "state " + cls); s.append(h("i", "", icon), document.createTextNode(label)); return s;
+}
+
+function renderCards(d) {
+  const cards = document.getElementById("cards");
+  cards.replaceChildren();
+  for (const key of ["ttft", "availability", "shed"]) {
+    const s = d.slo[key], c = h("div", "card");
+    const target = key === "ttft" ? `${pct(s.target, 0)} of requests get a first token within ${secs(s.threshold_s)}`
+                 : key === "availability" ? `${pct(s.target, 0)} of requests not failed by the server`
+                 : `at most ${pct(1 - s.target, 0)} of requests turned away (429) at rated load`;
+    c.append(h("h2", "", SLO_NAMES[key]), h("div", "line", "Target: " + target));
+    c.appendChild(h("div", "value", s.total ? pct(s.compliance, 2) : "–"));
+    c.appendChild(h("div", "line", s.total ? `${s.good.toLocaleString()} of ${s.total.toLocaleString()} requests` : "No requests yet"));
+    if (key === "shed") c.appendChild(state("none", "i", "Shown, not judged: the rated load for this server is not measured yet"));
+    else if (!s.total) c.appendChild(state("none", "–", "No data"));
+    else c.appendChild(s.met ? state("good", "✓", "Met") : state("bad", "✕", "Missed"));
+    if (s.total && key !== "shed") {
+      const left = Math.max(0, s.budget_left), mtr = h("div", "meter"), fill = h("div");
+      fill.style.width = (100 * left).toFixed(1) + "%"; mtr.appendChild(fill);
+      mtr.setAttribute("role", "meter"); mtr.setAttribute("aria-valuenow", (100 * left).toFixed(1));
+      mtr.setAttribute("aria-valuemin", "0"); mtr.setAttribute("aria-valuemax", "100");
+      mtr.setAttribute("aria-label", "Error budget left");
+      c.append(mtr, h("div", "line", s.budget_left > 0 ? `Error budget left: ${pct(s.budget_left, 0)}` : "Error budget used up"));
+    }
+    if (s.burn_rate_1h !== null) {
+      const line = h("div", "line", `Last hour: burn rate ${s.burn_rate_1h.toFixed(1)}× over ${s.requests_1h} requests `);
+      if (s.fast_burn && key !== "shed") line.appendChild(state("warn", "!", `fast burn (> ${d.fast_burn_threshold}×)`));
+      else if (s.requests_1h < d.burn_min_requests) line.appendChild(document.createTextNode(`(too few to flag, < ${d.burn_min_requests})`));
+      c.appendChild(line);
+    }
+    cards.appendChild(c);
+  }
+  document.getElementById("window").textContent =
+    `SLO window: rolling 7 days, counted since the last restart (${dur(d.slo_window_s)} so far). ` +
+    "Burn rate = the last hour's bad share ÷ the budgeted share; a fast burn above " +
+    `${d.fast_burn_threshold}× would spend ~9% of the 7-day budget in an hour.`;
+}
+
+function render(d) {
+  const v = Object.entries(d.variants)[0] || ["?", {}];
+  document.getElementById("meta").textContent =
+    `${v[1].version || v[0]} · norm ${v[1].norm_backend || "?"} · ${v[1].n_slots || "?"} slots · up ${dur(d.uptime_s)}`;
+  renderCards(d);
+  const mins = d.minutes.map((p) => p.covered ? { ...p, x: -p.minutes_ago } : null);
+  const threshold = d.slo.ttft.threshold_s;
+  document.getElementById("ref-label").textContent = `SLO: first token ≤ ${secs(threshold)}`;
+  drawChart(document.getElementById("c-req"), {
+    kind: "stack", points: mins, total: (p) => p ? p.ok + p.shed + p.failed : 0, ymin: 4, integer: true,
+    fmt: (v) => Number.isInteger(v) ? v.toLocaleString() : v.toFixed(1),
+    fmtTip: (v) => (v || 0).toLocaleString(), when: (p) => minLabel(p.minutes_ago),
+    series: [{ name: "Served", color: "var(--s1)", get: (p) => p.ok },
+             { name: "Busy (429)", color: "var(--s2)", get: (p) => p.shed },
+             { name: "Server error", color: "var(--s3)", get: (p) => p.failed }] });
+  drawChart(document.getElementById("c-lat"), {
+    kind: "line", points: mins, ymin: 0.1,
+    total: (p) => p ? Math.max(p.ttft_p95_s || 0, p.latency_p95_s || 0) : 0,
+    ref: { value: threshold, label: "SLO " + secs(threshold) },
+    fmt: (v) => v < 1 ? Math.round(v * 1000) + " ms" : (Number.isInteger(v) ? v : v.toFixed(1)) + " s",
+    fmtTip: secs, when: (p) => minLabel(p.minutes_ago),
+    series: [{ name: "First token p50", color: "var(--s1)", get: (p) => p.ttft_p50_s },
+             { name: "First token p95", color: "var(--s2)", get: (p) => p.ttft_p95_s },
+             { name: "Whole request p95", color: "var(--s3)", get: (p) => p.latency_p95_s }] });
+  const slots = v[1].n_slots || 1;
+  document.getElementById("slots-label").textContent = `${slots} slots`;
+  const snaps = d.snapshots.map((s) => ({ ...s, x: -s.seconds_ago / 60 }));
+  drawChart(document.getElementById("c-load"), {
+    kind: "line", points: snaps, maxGap: 0.5, ymin: slots, integer: true,
+    total: (p) => p ? Math.max(p.busy_slots, p.queue_len) : 0,
+    ref: { value: slots, label: `${slots} slots` },
+    fmt: (v) => Number.isInteger(v) ? String(v) : v.toFixed(1),
+    fmtTip: (v) => String(v), when: (p) => `${Math.round(p.seconds_ago)} s ago`,
+    series: [{ name: "Busy slots", color: "var(--s1)", get: (p) => p.busy_slots },
+             { name: "Waiting in queue", color: "var(--s2)", get: (p) => p.queue_len }] });
+  renderTable(d);
+}
+
+function renderTable(d) {
+  const t = document.getElementById("table");
+  t.replaceChildren();
+  const head = h("tr");
+  for (const c of ["Minute", "Requests", "Served", "429", "5xx", "First token p50", "p95", "Whole p95", "Tokens"])
+    head.appendChild(h("th", "", c));
+  t.appendChild(head);
+  for (const p of [...d.minutes].reverse()) {
+    if (!p.covered) continue;
+    const r = h("tr");
+    for (const c of [minLabel(p.minutes_ago), p.requests, p.ok, p.shed, p.failed, secs(p.ttft_p50_s),
+                     secs(p.ttft_p95_s), secs(p.latency_p95_s), p.completion_tokens.toLocaleString()])
+      r.appendChild(h("td", "", String(c)));
+    t.appendChild(r);
+  }
+}
+
+async function refresh() {
+  const banner = document.getElementById("banner");
+  try {
+    const r = await fetch("stats/history", { cache: "no-store" });
+    if (!r.ok) throw new Error("HTTP " + r.status);
+    last = await r.json();
+    banner.classList.remove("show");
+    document.querySelector("main").classList.remove("stale");
+    render(last);
+  } catch (e) {
+    banner.textContent = "Can't reach the server right now (it may be starting or asleep). Retrying every 5 s.";
+    banner.classList.add("show");
+    document.querySelector("main").classList.add("stale");
+  }
+}
+refresh();
+setInterval(refresh, 5000);
+let rt;
+window.addEventListener("resize", () => { clearTimeout(rt); rt = setTimeout(() => last && render(last), 150); });
+</script>
+</body>
+</html>