gta / summary_streaming_inference.py
Upload folder using huggingface_hub
ae4e95d verified 3 days ago
45.7 kB
"""Summaryformer streaming session: one trunk, one full-res stage, one tick per body frame.
One ``SummaryKVCache`` — one K/V buffer per layer,
``| Z0 | RAW : (Y+1) slots | SUM : (W−Y) slots |`` — one ``StageSpec``, model calls only at the
timesteps on the spec.
    z0 (GT, from the upload)  --forward_z0 t=0, note_z0_written-->  Z0 slot   (never generated)
    frame t (abs_t = 1, 2, ...), i frames already committed:
        T = SinkChaser.rotate(kv, model.rope_table, n_committed=i)   # one Z0 slot, one int (0 = none)
        kv.begin_frame()                                        # donate RAW[cursor_r]'s X summary rows into SUM
        lat ~ N(0,1) fp32
        for t in spec.timesteps (e.g. 999, 878, 669):
            v = forward_body(lat, t)  ->  x0_hat = signal(t) * lat - noise(t) * v   # writes record(t) into RAW[cursor_r]
            lat = signal(t_next) * x0_hat + noise(t_next) * z,  z FRESH             (re-noise to the next timestep)
        cache_mode "recache"          (KR): t=0 pass rewrites that RAW slot with the clean record, then commit
        cache_mode "last_step_writes" (KNR): commit the record the last pass wrote (no extra pass)
        commit = cursor_r += 1 (mod Y+1); no copy
    -> x0 (diag-gauss) -> un-gaussianize (rest stats) -> SAE decode_steady -> 4 RGB frames
Geometry is the model's: X summary tokens, Y raw-window frames, W summary-window frames, one
full-res grid (24x42 serve / the spec's ``hw``). One dense attention call per layer (every
query reads the prefix ``[0 : end]``); nothing on the cache path may ``torch.cat``.
Numerics: bf16 autocast on CUDA, fp32 latents / x0 arithmetic, bf16 KV (fp32 KV on CPU so
parity tests can be bit-exact). Weight dtype is the server's business.
torch.compile: the server hands in ``forward_body_compiled`` / ``decode_steady_compiled``.
``forward_body`` runs compiled from body frame 1 (frame id is a persistent 0-d int64 device
tensor filled in place; a tensor id selects ``plan_body(device_ends=True)``, whose ``end`` is
the one ``end_dev``). z0 warm is eager. SAE decode is compiled only once ``n_body >= Y`` unless
the flag says otherwise. KV storage (``k``, ``v``, the recency bias) and ``end_dev`` are pinned
(``mark_static_address``) after ``build`` and ``enable_recency``, so cudagraphs replay against
them. The wrapper marks those; this session does not mark them twice.
Long horizon: ``SinkChaser`` rotates the one Z0 slot when ``oldest_alive - sink_id`` reaches
``sink_gap_max`` (one int). Absolute ids keep growing; the session HARD-STOPS at the last RoPE
table row instead of wrapping.
Fail-fast: spec vs model at ctor; ``step()`` before ``start()``; abs_t >= RoPE rows;
a None sink gap (an unbounded stream needs the one slot chased).
Recency comes from ``StageSpec.recency`` and is enabled on the empty buffer right after build /
reset, before any body frame and before any compiled marking. The bias covers the whole buffer
and is non-zero only on the SUM region.
z0 focus region rides the z0 warm's cond ONLY, never body frames.
"""
from __future__ import annotations

import time
from collections import deque
from dataclasses import dataclass

import torch

from gta_streaming_kv_cache import SummaryKVCache
from inference_common import (
    StageSampling,
    cond_dict,
    frame_tokens,
    warm_z0_clean,
)
from sae import diag_gaussian_latent_to_training_latent, training_latent_to_diag_gaussian_latent
from summaryformer import Summaryformer
from summaryformer_utils.streaming_sink import (
    SinkChaser,
    sink_gaps_describe,
    sink_gaps_from_config,
)
from summaryformer_wrapper import SummaryformerWrapper

# How a generated frame's keys and values get into the cache: "recache" runs one extra clean pass over
# the finished frame; "last_step_writes" keeps what the final denoising step wrote. The checkpoint's
# config names which one it expects.
CACHE_MODES = ("recache", "last_step_writes")

Z0_C = Z_REST_C = int(Summaryformer.LATENT_CHANNELS)
SAE_SPATIAL = 16
_RGB_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1)
_RGB_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1)


@dataclass(frozen=True)
class StageSpec:
    """ONE stage; the server builds it from the ckpt config and the session checks it against the model.
    ``samp``: ``resolve_sampling(ckpt_cfg)``, the timesteps a body frame is denoised at.
    ``stats_*``: full-res diag-gauss artifacts (C=48). ``sink_gap_max``: one int >= 2 (the one Z0 slot).
    """
    hw: tuple[int, int]
    stats_z0: dict
    stats_rest: dict
    samp: StageSampling
    cache_mode: str
    sink_gap_max: int
    recency: dict | None
    focus_region: torch.Tensor | None

    def __post_init__(self):
        h, w = (int(self.hw[0]), int(self.hw[1]))
        if h <= 0 or w <= 0:
            raise ValueError(f"StageSpec.hw must be positive (h, w); got {self.hw!r}")
        if not isinstance(self.samp, StageSampling):
            raise TypeError(f"StageSpec.samp must be StageSampling; got {type(self.samp)}")
        if len(self.samp.timesteps) < 1:
            raise ValueError(f"StageSpec.samp has no model timesteps: {self.samp.timesteps}")
        if self.cache_mode not in CACHE_MODES:
            raise ValueError(f"StageSpec.cache_mode={self.cache_mode!r} not in {CACHE_MODES}")
        gap = sink_gaps_from_config(self.sink_gap_max)
        if gap is None:
            raise RuntimeError(
                "StageSpec.sink_gap_max is None: an unbounded stream needs the one Z0 slot chased "
                f"({sink_gaps_describe(gap)}; pass an int >= 2)"
            )
        object.__setattr__(self, "sink_gap_max", gap)
        if self.recency is not None and not isinstance(self.recency, dict):
            raise TypeError(f"StageSpec.recency must be a dict or None; got {type(self.recency)}")
        for name, art in (("stats_z0", self.stats_z0), ("stats_rest", self.stats_rest)):
            if not (isinstance(art, dict) and "mean" in art and "std" in art):
                raise ValueError(f"StageSpec.{name} must be a diag-gauss dict with mean/std")
            if int(torch.as_tensor(art["mean"]).numel()) != Z0_C:
                raise ValueError(
                    f"StageSpec.{name} mean numel={int(torch.as_tensor(art['mean']).numel())} != C={Z0_C}"
                )
        region = self.focus_region
        if region is None:
            return
        if not (
            isinstance(region, torch.Tensor)
            and region.dtype == torch.bool
            and tuple(region.shape) == (1, h, w)
        ):
            raise RuntimeError(
                f"StageSpec.focus_region must be Bool [1,{h},{w}] or None; "
                f"got {getattr(region, 'shape', type(region))} dtype={getattr(region, 'dtype', None)}"
            )
        if not bool(region.any()):
            raise RuntimeError("StageSpec.focus_region is all-False: pass None for 'no region'")

    @property
    def h(self) -> int:
        return int(self.hw[0])

    @property
    def w(self) -> int:
        return int(self.hw[1])

    @property
    def timesteps(self) -> list[int]:
        """Timesteps the model is called at on body frames, high -> low (e.g. [999, 878, 669]). The session
        only ever denoises body frames; the seed frame is written clean at t = 0."""
        return list(self.samp.timesteps)

    @property
    def spaced_indices(self) -> list[int]:
        """Indices into the sampling grid's tables the loop visits, in order, len == nfe
        (repeats when the samp carries a step_plan)."""
        return list(self.samp.indices)

    @property
    def nfe(self) -> int:
        return len(self.timesteps)


def _stats_to_device(stats: dict, device: torch.device) -> dict:
    """Diag-gauss artifact with mean/std already on `device` (fp32): no H2D per call."""
    out = dict(stats)
    out["mean"] = torch.as_tensor(stats["mean"]).to(device=device, dtype=torch.float32)
    out["std"] = torch.as_tensor(stats["std"]).to(device=device, dtype=torch.float32)
    return out


def _check_spec_against_model(model, spec: StageSpec) -> None:
    """X/Y/W live on the model."""
    x, y, w = (
        int(model.summary_tokens),
        int(model.raw_window_frames),
        int(model.summary_window_frames),
    )
    if x < 1 or y < 0 or w < 0:
        raise RuntimeError(f"[summary-stream] model X/Y/W={(x, y, w)} is not a serveable geometry")


# ``start(focus_region=FOCUS_DEFAULT)``: use the spec's region (the ckpt's served default, a centred diamond when the
# focus operator was trained). ``None`` = no region for this game; a Bool ``[1, h, w]`` = a custom region (the client's
# painted mask) for this game only. The default is restored by the next ``start`` that does not override it.
FOCUS_DEFAULT = object()


def _coerce_focus_region(region, spec: StageSpec, device) -> torch.Tensor | None:
    if region is None:
        return None
    if not (
        isinstance(region, torch.Tensor)
        and region.dtype == torch.bool
        and tuple(region.shape) == (1, spec.h, spec.w)
    ):
        raise RuntimeError(
            f"[summary-stream] focus_region must be Bool [1,{spec.h},{spec.w}] or None; "
            f"got {getattr(region, 'shape', type(region))} dtype={getattr(region, 'dtype', None)}"
        )
    if not bool(region.any()):
        raise RuntimeError("[summary-stream] focus_region is all-False: pass None for 'no region'")
    return region.to(device=device).contiguous()


def _normalize_fps(fps) -> float | None:
    """NaN / None = no fps conditioning; a number in [1, 120] is the game clock."""
    if fps is None:
        return None
    if isinstance(fps, torch.Tensor):
        if fps.numel() != 1:
            raise RuntimeError(f"[summary-stream] fps tensor must be a scalar; got {tuple(fps.shape)}")
        fps = float(fps.item())
    if isinstance(fps, (int, float)):
        f = float(fps)
        if f != f:  # NaN
            return None
        if not 1.0 <= f <= 120.0:
            raise RuntimeError(f"[summary-stream] fps must be None/NaN or a number in [1, 120]; got {fps!r}")
        return f
    raise RuntimeError(f"[summary-stream] fps must be None/NaN or a number in [1, 120]; got {fps!r}")


# ----------------------------------------------------------------------------------------------
# SAE sphere projection (serve-time option)
# ----------------------------------------------------------------------------------------------
# The SAE encoder ends in `sae.SphereNorm3d`: every 48-channel token is forced to RMS=1 (L2 radius
# sqrt(48)), no mean subtraction, non-affine. The decoder expects such tokens and has no input norm of
# its own. The diffusion model works in per-channel diag-gauss space, where that sphere is an
# ellipsoid, so a generated x0_hat can drift off the sphere. These modes put it back:
#   "off"    nothing.
#   "decode" project the RAW latent handed to the SAE decoder only (seed and every body frame). The
#            diffusion loop and the KV are untouched.
#   "all"    "decode" plus: every x0_hat inside the renoise loop (ungauss -> project -> regauss, so
#            the next re-noise, the t=0 recache write and the decode all see the on-sphere x0), and
#            the seed z0 in start() (a no-op for an encoder z0, which is already on the sphere).
SPHERE_MODES = ("off", "decode", "all")
_SPHERE_EPS = 1e-6   # sae.SphereNorm3d: eps added INSIDE the sqrt


def sphere_project_raw(z_raw: torch.Tensor, *, cdim: int = 1) -> torch.Tensor:
    """RAW SAE latent -> the encoder's SphereNorm3d manifold: per-token RMS=1 over the channel axis.
    Bit-for-bit the encoder's op (fp32 rms, eps inside the sqrt). Shape-agnostic; ``cdim`` is the
    48-channel axis."""
    rms = z_raw.float().pow(2).mean(dim=cdim, keepdim=True).add(_SPHERE_EPS).sqrt()
    return z_raw / rms.to(z_raw.dtype)


def _sched_tables(samp: StageSampling, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    """The sampling grid's signal / noise scales, sqrt(ā) / sqrt(1-ā), on `device` (fp32), indexed by grid
    index."""
    a = torch.as_tensor(samp.grid.signal, device=device, dtype=torch.float32)
    b = torch.as_tensor(samp.grid.noise, device=device, dtype=torch.float32)
    return a, b


def _body_steps(samp: StageSampling) -> tuple[list[int], list[int]]:
    """(grid indices, timesteps) the body loop visits, in order, len == NFE. Each grid index once
    high -> low, or ``samp.step_plan`` (repeats allowed: a repeated index re-noises x0_hat to the same t
    and calls the model there again)."""
    return list(samp.indices), list(samp.timesteps)


def renoise_body_x0(
    *,
    call,
    cond: dict,
    samp: StageSampling,
    lat_shape: tuple[int, ...],
    noise_fn,
    sqrt_ab: torch.Tensor,
    sqrt_om: torch.Tensor,
    t_buffer: torch.Tensor,
    project_x0=None,
) -> torch.Tensor:
    """K renoise passes on one body frame. Model at REAL timesteps; fresh z between passes.
    ``call(x[B,1,C,H,W], t[B,1], cond) -> v[B,1,C,H,W]``. ``noise_fn(shape) -> tensor``.
    ``sqrt_ab`` / ``sqrt_om`` are the device-resident signal / noise tables of the sampling grid
    (``_sched_tables``). Returns x0
    ``[B,1,C,H,W]`` fp32. ``project_x0`` (optional, eager, fp32 gauss -> fp32 gauss) is applied to
    every x0 estimate before it is re-noised or returned (sphere mode "all").
    """
    indices, real_ts = _body_steps(samp)
    if len(indices) != len(real_ts) or not indices:
        raise RuntimeError(f"[summary-stream] spaced/real timestep mismatch: {indices} vs {real_ts}")
    lat = noise_fn(lat_shape)
    if not (isinstance(lat, torch.Tensor) and lat.dim() == 5 and int(lat.size(1)) == 1):
        raise RuntimeError(f"[summary-stream] noise_fn must return [B,1,C,H,W]; got {getattr(lat, 'shape', type(lat))}")
    x0 = None
    for step_i, (t_idx, t_real) in enumerate(zip(indices, real_ts)):
        t_buffer.fill_(float(t_real))
        v = call(lat, t_buffer, cond).float()
        a, b = sqrt_ab[int(t_idx)], sqrt_om[int(t_idx)]
        x0 = (a * lat - b * v) / (a * a + b * b)
        if project_x0 is not None:
            x0 = project_x0(x0)
        if step_i < len(indices) - 1:
            z = noise_fn(tuple(x0.shape))
            a_n, b_n = sqrt_ab[int(indices[step_i + 1])], sqrt_om[int(indices[step_i + 1])]
            lat = a_n * x0 + b_n * z
    return x0


class SummaryStreamSession:
    """One player, one device. Owns the KV buffer, RNG, SAE decode cache, gaussianization boundary."""

    def __init__(
        self,
        *,
        generator,
        sae_decoder,
        spec: StageSpec,
        device,
        rgb_hw: tuple[int, int] = (384, 672),
        forward_body_compiled=None,
        decode_steady_compiled=None,
        use_compiled_sae_in_steady_state: bool = True,
    ):
        gen = getattr(generator, "module", generator)
        if not isinstance(gen, SummaryformerWrapper):
            raise RuntimeError(f"[summary-stream] generator must be SummaryformerWrapper; got {type(generator)}")
        if not isinstance(spec, StageSpec):
            raise TypeError(f"[summary-stream] spec must be StageSpec; got {type(spec)}")
        if not (hasattr(sae_decoder, "decode_bootstrap") and hasattr(sae_decoder, "decode_steady")):
            raise RuntimeError("[summary-stream] sae_decoder must expose decode_bootstrap / decode_steady")
        device = torch.device(device)
        model = gen.model
        _check_spec_against_model(model, spec)
        rh, rw = (int(x) for x in rgb_hw)
        if (rh, rw) != (spec.h * SAE_SPATIAL, spec.w * SAE_SPATIAL):
            raise RuntimeError(
                f"[summary-stream] rgb_hw {rh}x{rw} != {SAE_SPATIAL} * latent grid {spec.h}x{spec.w}"
            )

        self.generator, self.model, self.sae, self.device = gen, model, sae_decoder, device
        self.spec = spec
        self.focus_region_default = self.focus_region = _coerce_focus_region(spec.focus_region, spec, device)
        self.rgb_hw = (rh, rw)
        # Serve: bf16 autocast + bf16 KV. CPU parity: fp32 throughout so x0 can be torch.equal.
        self.amp_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        try:
            self.sae_dtype = next(sae_decoder.parameters()).dtype
        except StopIteration:
            self.sae_dtype = torch.float32
        self.tf = int(frame_tokens(model, spec.h, spec.w))
        self.X = int(model.summary_tokens)
        self.Y = int(model.raw_window_frames)
        self.W = int(model.summary_window_frames)

        self.visual_kv: SummaryKVCache = gen.new_visual_kv_cache()
        if not isinstance(self.visual_kv, SummaryKVCache):
            raise RuntimeError("[summary-stream] new_visual_kv_cache must return SummaryKVCache")
        with torch.no_grad():
            self.visual_kv.build(
                batch=1, frame_tokens=self.tf, dtype=self.amp_dtype, device=device,
            )
            # Recency on the empty buffer (bias non-zero only on SUM), before any frame and before
            # static marking. mark_streaming_caches_static pins k, v, the bias, and the one end_dev.
            self.visual_kv.enable_recency(spec.recency)
        self.recency = self.visual_kv.recency
        self._chaser = self._new_chaser()

        self._body_x, self._sae_x = forward_body_compiled, decode_steady_compiled
        self.use_compiled_sae_in_steady_state = bool(use_compiled_sae_in_steady_state)
        self._mode_log: dict[str, str] = {}
        self._any_generator_compiled = self._body_x is not None
        if self._any_generator_compiled:
            n_static = int(SummaryformerWrapper.mark_streaming_caches_static(self.visual_kv))
            if n_static <= 0:
                raise RuntimeError("[summary-stream] compiled generator given but no KV tensors were marked static")
            print(f"[perf][cudagraph] marked {n_static} KV tensors static (mark_static_address) at build", flush=True)

        self._rng: torch.Generator | None = None
        self.noise_fn = None  # optional ``(shape) -> tensor`` hook (parity tests)
        self._abs_t: int | None = None
        self._abs_t_dev = torch.zeros((), device=device, dtype=torch.long)
        self._n_body = 0
        self._sae_cache = None
        self.last_x0: torch.Tensor | None = None
        self.last_ms: dict[str, float] = {}
        self._fps_logged: float | None | str = "unset"
        self._peak_ema: float | None = None

        self._stats_dev = {
            "z0": _stats_to_device(spec.stats_z0, device),
            "rest": _stats_to_device(spec.stats_rest, device),
        }
        self._sqrt_ab, self._sqrt_om = _sched_tables(spec.samp, device)
        # Forward re-prefill history (see reprefill()): the last W committed body frames as
        # (gauss x0 [1,C,h,w] fp32 AS WRITTEN by the recache pass, action [1,1,5], fps), plus the
        # game's gauss z0. ~190 KB/frame at 24x42; 6 MB at W=32.
        self._hist: deque = deque(maxlen=int(self.W))
        self._z0_gauss: torch.Tensor | None = None
        self._z0_fps: float | None = None
        self.reprefill_every: int = 0   # 0 = off; n = rebuild the KV from history every n body frames
        self._since_reprefill = 0       # body frames since start() / the last reprefill (the trigger counter)
        self.n_reprefills = 0           # rebuilds this game
        self.last_reprefill: dict | None = None
        self._t_dev = torch.zeros((1, 1), device=device, dtype=torch.float32)
        self._fps_dev = torch.zeros((1,), device=device, dtype=torch.float32)
        self._action_dev = torch.zeros((1, 1, 5), device=device, dtype=torch.float32)
        self._action_pin = (
            torch.zeros((1, 1, 5), dtype=torch.float32).pin_memory() if device.type == "cuda" else None
        )
        self._rgb_mean, self._rgb_std = _RGB_MEAN.to(device), _RGB_STD.to(device)
        self._ev: dict[str, tuple] = {}
        if self._any_generator_compiled:
            for t_ in (self._abs_t_dev, self._t_dev, self._fps_dev, self._action_dev):
                torch._dynamo.mark_static_address(t_)
        self._print_rings()

    def _new_chaser(self) -> SinkChaser:
        return SinkChaser(
            raw_window=self.Y, summary_window=self.W, gaps=self.spec.sink_gap_max, first_id=1,
        )

    def _print_rings(self) -> None:
        kv = self.visual_kv
        rec = self.recency
        print(
            f"[summary-stream] one trunk X={self.X} Y={self.Y} W={self.W} Tf={self.tf} "
            f"grid={self.spec.h}x{self.spec.w} rgb={self.rgb_hw[0]}x{self.rgb_hw[1]} "
            f"cache_mode={self.spec.cache_mode} amp={self.amp_dtype} "
            f"{sink_gaps_describe(self.spec.sink_gap_max)} "
            f"rope_rows={int(self.model.rope_table.size(0))} (hard stop) "
            f"compiled(body={self._body_x is not None}, sae={self._sae_x is not None}) "
            f"body compiled from frame 1 (device frame id); sae from steady state",
            flush=True,
        )
        print(f"[summary-stream] {kv.describe()}  model t={self.spec.timesteps} nfe={self.spec.nfe}", flush=True)
        print(
            f"[summary-stream] {'recency OFF' if rec is None else rec.cfg.describe()}",
            flush=True,
        )
        if rec is not None:
            from summaryformer_utils.recency_bias import prepare_flex_decode
            # One prefix read: queries are the whole record (Tf + X), keys are the whole buffer.
            print(
                f"[summary-stream] biased attend: {prepare_flex_decode(cap=int(kv.capacity), batch=1, kv_heads=int(self.model.num_heads), frame_tokens=self.tf + self.X, device=self.device)}",
                flush=True,
            )
        print(f"[summary-stream] {self._focus_region_desc()}", flush=True)

    def _focus_region_desc(self) -> str:
        r = self.focus_region
        if r is None:
            return "z0 focus region OFF (no marker on z0)"
        return (
            f"z0 focus region ON: {int(r.sum().item())}/{int(r.numel())} z0 tokens marked "
            f"({100.0 * float(r.float().mean()):.0f}%), z0 warm only"
        )

    @property
    def n_body(self) -> int:
        return int(self._n_body)

    @property
    def abs_t(self) -> int | None:
        return self._abs_t

    @property
    def chaser(self) -> SinkChaser:
        return self._chaser

    def _gauss_in(self, z_raw: torch.Tensor, *, which: str) -> torch.Tensor:
        return training_latent_to_diag_gaussian_latent(z_raw.float(), self._stats_dev[which])

    def _gauss_out(self, z: torch.Tensor, *, which: str) -> torch.Tensor:
        """Gauss -> RAW for the SAE decoder. Sphere modes "decode" / "all" project the RAW token
        onto the encoder's RMS=1 sphere here, in fp32, before the dtype cast and the compiled decode."""
        z_raw = diag_gaussian_latent_to_training_latent(z.float(), self._stats_dev[which])
        if self.sphere_mode != "off":
            z_raw = sphere_project_raw(z_raw, cdim=1)
        return z_raw.to(dtype=self.sae_dtype).unsqueeze(2).contiguous(memory_format=torch.channels_last_3d)

    # Sphere projection mode (SPHERE_MODES). Plain attribute, read per call; no compiled graph depends
    # on it (every projection is an eager fp32 op on the tensor BEFORE a compiled call).
    sphere_mode: str = "off"

    def set_sphere_mode(self, mode: str) -> None:
        if mode not in SPHERE_MODES:
            raise ValueError(f"[summary-stream] sphere mode {mode!r} not in {SPHERE_MODES}")
        if mode != self.sphere_mode:
            print(f"[summary-stream] sphere projection {self.sphere_mode} -> {mode}", flush=True)
        self.sphere_mode = mode

    def _project_gauss(self, z: torch.Tensor, *, which: str) -> torch.Tensor:
        """Gauss -> RAW -> sphere -> gauss, fp32. ``z`` is the [B,1,C,H,W] renoise x0 (the diag-gauss
        helpers want C at dim 1, so the T=1 axis is squeezed around them)."""
        stats = self._stats_dev[which]
        five = z.dim() == 5
        zb = z.float().squeeze(1) if five else z.float()
        z_raw = sphere_project_raw(diag_gaussian_latent_to_training_latent(zb, stats), cdim=1)
        out = training_latent_to_diag_gaussian_latent(z_raw, stats)
        return out.unsqueeze(1) if five else out

    def _coerce_bchw(self, z, *, name: str) -> torch.Tensor:
        if isinstance(z, torch.Tensor) and z.dim() == 5 and 1 in (int(z.size(1)), int(z.size(2))):
            z = z.squeeze(1) if int(z.size(1)) == 1 else z.squeeze(2)
        if not (isinstance(z, torch.Tensor) and tuple(z.shape) == (1, Z0_C, self.spec.h, self.spec.w)):
            raise RuntimeError(
                f"[summary-stream] {name} must be [1,{Z0_C},{self.spec.h},{self.spec.w}]; "
                f"got {getattr(z, 'shape', type(z))}"
            )
        return z.to(device=self.device).contiguous()

    def _rgb_out(self, rgb_btchw: torch.Tensor, *, frames: int) -> torch.Tensor:
        rh, rw = self.rgb_hw
        if tuple(rgb_btchw.shape) != (1, frames, 3, rh, rw):
            raise RuntimeError(
                f"[summary-stream] SAE rgb {tuple(rgb_btchw.shape)} != (1,{frames},3,{rh},{rw})"
            )
        return (rgb_btchw.float() * self._rgb_std + self._rgb_mean).clamp(0.0, 1.0)

    def _autocast(self, dtype: torch.dtype | None = None):
        dt = self.amp_dtype if dtype is None else dtype
        return torch.autocast(
            device_type="cuda", dtype=dt, enabled=self.device.type == "cuda" and dt != torch.float32,
        )

    def _timed(self, name: str, fn):
        if self.device.type == "cuda":
            s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
            s.record()
            out = fn()
            e.record()
            self._ev[name] = (s, e)
            return out
        t0 = time.perf_counter()
        out = fn()
        self._ev[name] = (t0, time.perf_counter())
        return out

    def _collect_ms(self) -> dict[str, float]:
        ms = {
            k: float(s.elapsed_time(e)) if self.device.type == "cuda" else (e - s) * 1000.0
            for k, (s, e) in self._ev.items()
        }
        self._ev.clear()
        return ms

    def _in_steady_state(self) -> bool:
        return self._n_body >= int(self.Y)

    def _pick(self, what: str, eager, compiled, use_compiled: bool):
        mode = "compiled" if (compiled is not None and use_compiled) else "eager"
        if self._mode_log.get(what) != mode:
            print(f"[perf][{what}-mode] mode={mode} abs_t={self._abs_t} n_body={self._n_body}", flush=True)
            self._mode_log[what] = mode
        return compiled if mode == "compiled" else eager

    def _forward_body(self):
        return self._pick("generator-body", self.generator.forward_body, self._body_x, True)

    def _decode_steady(self):
        return self._pick(
            "sae-decode", self.sae.decode_steady, self._sae_x,
            self._in_steady_state() or not self.use_compiled_sae_in_steady_state,
        )

    def _cond(self, fps: float | None, action_b15: torch.Tensor | None = None, **z0_knobs) -> dict:
        fps_v = None
        if fps is not None:
            self._fps_dev.fill_(float(fps))
            fps_v = self._fps_dev
        return cond_dict(
            fps=fps_v,
            action_context=action_b15,
            **z0_knobs,
        )

    def _action_to_device(self, action_b15: torch.Tensor) -> torch.Tensor:
        if action_b15.device == self._action_dev.device:
            self._action_dev.copy_(action_b15.to(torch.float32))
        elif self._action_pin is not None:
            self._action_pin.copy_(action_b15.to(torch.float32))
            self._action_dev.copy_(self._action_pin, non_blocking=True)
        else:
            self._action_dev.copy_(action_b15.to(self._action_dev.device, torch.float32))
        return self._action_dev

    def _noise(self, shape) -> torch.Tensor:
        if self.noise_fn is not None:
            z = self.noise_fn(shape)
            return z.to(device=self.device, dtype=torch.float32)
        if self._rng is None:
            raise RuntimeError("[summary-stream] RNG missing (start() first) and no noise_fn")
        return torch.randn(shape, generator=self._rng, device=self.device, dtype=torch.float32)

    def _frame_id_arg(self, abs_t: int, *, compiled: bool):
        """Compiled form: persistent 0-d int64 filled in place. The model treats a tensor id as
        ``plan_body(device_ends=True)`` and reads the one ``end_dev``. Eager: the Python int."""
        if compiled:
            self._abs_t_dev.fill_(int(abs_t))
            return self._abs_t_dev
        return int(abs_t)

    def _body_call(self, x: torch.Tensor, t_buf: torch.Tensor, cond: dict, *, abs_t: int):
        fn = self._forward_body()
        compiled = fn is self._body_x
        with self._autocast():
            v = fn(
                cond, x.to(self.amp_dtype), t_buf, kv_cache=self.visual_kv,
                absolute_start_frames=self._frame_id_arg(abs_t, compiled=compiled),
            )
        return v

    def _denoise_frame(self, cond: dict, *, abs_t: int) -> torch.Tensor:
        """One body frame: the denoising passes, then the cache write the ckpt's cache_mode asks for. B=1.
        ``begin_frame`` once before the passes (the model's first ``forward_body`` calls it too;
        the second call is a no-op until commit). Donation of the RAW slot about to be overwritten
        happens here, before any write. Commit advances ``cursor_r`` and does not copy.
        """
        self.visual_kv.begin_frame()
        lat_shape = (1, 1, Z_REST_C, self.spec.h, self.spec.w)
        x0 = renoise_body_x0(
            call=lambda x, t, c, _abs=abs_t: self._body_call(x, t, c, abs_t=_abs),
            cond=cond,
            samp=self.spec.samp,
            lat_shape=lat_shape,
            noise_fn=self._noise,
            sqrt_ab=self._sqrt_ab,
            sqrt_om=self._sqrt_om,
            t_buffer=self._t_dev,
            # Sphere mode "all": every x0 estimate back onto the SAE sphere (eager fp32, before the
            # re-noise, the t=0 recache below, and the decode). "decode" leaves the loop alone.
            project_x0=(lambda x: self._project_gauss(x, which="rest")) if self.sphere_mode == "all" else None,
        )
        if self.spec.cache_mode == "recache":
            self._t_dev.fill_(0.0)
            self._body_call(x0, self._t_dev, cond, abs_t=abs_t)
            self.visual_kv.commit_current_frame()
            # What the cache now holds for this frame is the clean record of THIS x0: remember it.
            self._hist.append((
                x0[:, 0].detach().float().clone(),
                self._action_dev.detach().clone(),
                self._fps_logged if isinstance(self._fps_logged, (int, float)) else None,
            ))
        else:
            self.visual_kv.commit_current_frame()
        return x0

    # ------------------------------------------------------------------------------------------
    # RE-PREFILL: rebuild the whole KV from the clean x0 history
    # ------------------------------------------------------------------------------------------
    # After thousands of frames the buffer is not what a fresh forward over the same frames would
    # produce: Z0's K has been rotated in place in bf16 every (gap_max-1) frames (SinkChaser), and
    # every SUM row is a memcpy of K/V that were computed against whatever the prefix was when that
    # frame ran. reprefill() throws the buffer away and replays [z0 | last W clean x0]: the z0 warm,
    # then every stored frame as a clean t=0 write + commit. Ids are renumbered to 0..n, so the sink
    # gap is back to 1 and abs_t stays bounded (no RoPE-table hard stop). Only under "recache": in
    # "last_step_writes" the committed keys are noisy-input keys and a clean rebuild would be a
    # different key kind.

    @torch.no_grad()
    def reprefill(self) -> dict:
        if self._abs_t is None or self._z0_gauss is None:
            raise RuntimeError("[summary-stream] reprefill() before start()")
        if self.spec.cache_mode != "recache":
            raise RuntimeError(f"[summary-stream] reprefill needs cache_mode='recache'; got {self.spec.cache_mode!r}")
        hist = list(self._hist)
        n = len(hist)
        if n > int(self.W):
            raise RuntimeError(f"[summary-stream] history {n} > W={self.W}")
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        t0 = time.perf_counter()
        kv = self.visual_kv
        if bool(getattr(kv, "_body_open", False)):
            raise RuntimeError("[summary-stream] reprefill with an open (uncommitted) frame")
        old_abs_t, old_rot = int(self._abs_t), self._chaser.describe()

        # 1. Empty the buffer in place (same static tensors), recency on the empty buffer, fresh chaser.
        kv.reset()
        kv.enable_recency(self.spec.recency)
        self.recency = kv.recency
        self._chaser = self._new_chaser()

        # 2. Replay the history into the empty buffer.
        self._reprefill_sequential(hist)

        # 3. Counters follow the rebuilt buffer.
        self._n_body, self._abs_t, self._since_reprefill = n, n + 1, 0
        self.n_reprefills += 1
        expect_r = min(n, int(self.Y) + 1)
        expect_s = min(max(0, n - (int(self.Y) + 1)), max(0, int(self.W) - int(self.Y)))
        if int(kv.filled_r) != expect_r or int(kv.filled_s) != expect_s or int(kv.n_committed) != n or not kv.z0_written:
            raise RuntimeError(
                f"[summary-stream] reprefill bookkeeping: filled_r={kv.filled_r} (want {expect_r}) "
                f"filled_s={kv.filled_s} (want {expect_s}) n_committed={kv.n_committed} (want {n}) z0_written={kv.z0_written}"
            )
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        ms = (time.perf_counter() - t0) * 1e3
        self.last_reprefill = {"frames": n, "ms": ms, "abs_t_before": old_abs_t, "abs_t_after": n + 1}
        print(
            f"[reprefill] rebuilt KV from z0 + {n} clean frames in {ms:.0f} ms | abs_t {old_abs_t} -> {n + 1} | "
            f"{old_rot} -> {self._chaser.describe()} | {kv.describe()}",
            flush=True,
        )
        return self.last_reprefill

    def _reprefill_sequential(self, hist: list) -> None:
        """z0 clean warm at id 0, then every stored frame as the clean t=0 write the recache pass makes (ids 1..n),
        begin_frame / commit around each. n <= W so the sink never rotates inside the replay."""
        kv = self.visual_kv
        cond0 = self._cond(self._z0_fps, z0_focus_region=self.focus_region)
        with self._autocast():
            warm_z0_clean(self.generator, kv, cond0, self._z0_gauss.unsqueeze(1))
        for i, (x0, action, fps) in enumerate(hist):
            if self._chaser.rotation_before_frame(i):
                raise RuntimeError(f"[summary-stream] reprefill: sink would rotate at replay frame {i} (n={len(hist)}, W={self.W})")
            cond = self._cond(fps, self._action_to_device(action))
            kv.begin_frame()
            self._t_dev.fill_(0.0)
            self._body_call(x0.unsqueeze(1), self._t_dev, cond, abs_t=i + 1)
            kv.commit_current_frame()

    def _chase_sink(self) -> None:
        """Rotate the one Z0 slot before body frame i (``n_committed`` = frames already committed).
        ``T = chaser.rotate(kv, freqs, i)`` with ``i`` the number of body frames committed so far.
        ``T == 0`` means the gap has not reached ``sink_gap_max``.
        """
        T = self._chaser.rotate(self.visual_kv, self.model.rope_table, self._n_body)
        if T:
            print(f"[sink] T={T} ({self._chaser.describe()} n_body={self._n_body})", flush=True)

    @torch.no_grad()
    def start(self, z0_raw_bchw, *, fps, focus_region=FOCUS_DEFAULT) -> torch.Tensor:
        """New game. GT-warm z0 (t=0 write, never generated) under this fps clock, SAE bootstrap
        -> rgb0 ``[1,1,3,H,W]`` in ``[0,1]``. Next abs_t = 1.
        ``focus_region``: ``FOCUS_DEFAULT`` = the ckpt's served region (spec); ``None`` = none this game; a Bool
        ``[1, h, w]`` = the client's custom region for this game (the z0 warm AND every reprefill of it)."""
        z0_raw = self._coerce_bchw(z0_raw_bchw, name="z0_raw")
        if focus_region is FOCUS_DEFAULT:
            self.focus_region = self.focus_region_default
        else:
            self.focus_region = _coerce_focus_region(focus_region, self.spec, self.device)
        if self.sphere_mode == "all":
            # A first frame that is not encoder output is put on the sphere before the KV
            # warm. For an encoder z0 this is the identity up to eps.
            z0_raw = sphere_project_raw(z0_raw.float(), cdim=1)
        fps = _normalize_fps(fps)
        self.reset()
        self._rng = torch.Generator(device=self.device)
        self._rng.manual_seed(int(torch.randint(0, 2**31 - 1, (1,)).item()))

        cond = self._cond(fps, z0_focus_region=self.focus_region)
        self._fps_logged = fps
        z0 = self._gauss_in(z0_raw, which="z0")
        with self._autocast():
            warm_z0_clean(self.generator, self.visual_kv, cond, z0.unsqueeze(1))
        with self._autocast(self.sae_dtype):
            rgb0, self._sae_cache = self.sae.decode_bootstrap(self._gauss_out(z0, which="z0"))
        self.last_x0, self._abs_t, self._n_body = z0, 1, 0
        self._z0_gauss, self._z0_fps = z0.detach().float().clone(), fps
        self._hist.clear()
        print(
            f"[kv-boundary] z0 warmed (GT, t=0, fps={fps}) {self.visual_kv.describe()} "
            f"z0 mean={float(z0.mean()):.3f} std={float(z0.std()):.3f} "
            f"focus_region={'off' if self.focus_region is None else f'{int(self.focus_region.sum().item())} tokens'}",
            flush=True,
        )
        return self._rgb_out(rgb0, frames=1)

    @torch.no_grad()
    def step(self, action_b15, *, fps) -> torch.Tensor:
        """One body latent frame under this fps clock (None/NaN = no clock) + action -> rgb4 ``[1,4,3,H,W]``."""
        if self._abs_t is None or self._sae_cache is None:
            raise RuntimeError("[summary-stream] step() before start()")
        if not (isinstance(action_b15, torch.Tensor) and tuple(action_b15.shape) == (1, 1, 5)):
            raise RuntimeError(f"[summary-stream] action_b15 must be [1,1,5]; got {getattr(action_b15, 'shape', None)}")
        fps = _normalize_fps(fps)
        every = int(self.reprefill_every or 0)
        if every > 0 and self._since_reprefill >= every:
            self.reprefill()   # renumbers abs_t: read it after
        self._since_reprefill += 1
        t = int(self._abs_t)
        if t >= int(self.model.rope_table.size(0)):
            raise RuntimeError(
                f"[summary-stream] HARD STOP: abs_t={t} >= RoPE table rows {int(self.model.rope_table.size(0))}; "
                f"end the game, do not wrap"
            )
        self._chase_sink()
        cond = self._cond(fps, self._action_to_device(action_b15))
        if fps != self._fps_logged:
            print(f"[fps] clock {self._fps_logged} -> {fps} from abs_t={t}", flush=True)
            self._fps_logged = fps

        x0 = self._timed("model", lambda: self._denoise_frame(cond, abs_t=t))
        with self._autocast(self.sae_dtype):
            rgb4, self._sae_cache = self._timed(
                "sae", lambda: self._decode_steady()(self._gauss_out(x0[:, 0], which="rest"), self._sae_cache),
            )
        xf = x0[:, 0].float()
        stats_dev = torch.stack([
            xf.mean(), xf.std(), xf.abs().amax(),
            torch.isfinite(xf).all().float(), torch.isnan(xf).sum().float(),
        ])
        x_mean, x_std, x_absmax, x_finite, x_nan = stats_dev.tolist()
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        ms = self._collect_ms()
        if x_finite < 1.0:
            raise RuntimeError(f"[summary-stream] x0 at abs_t={t} is not finite (nan={int(x_nan)})")

        self._n_body += 1
        # Closed frame (commit already ran; the next begin_frame has not). RAW holds min(n, Y+1)
        # records. SUM receives a donation only once RAW is full, at the start of the frame that
        # overwrites the oldest slot, so after n commits it holds min(max(0, n-(Y+1)), W-Y).
        n = int(self._n_body)
        expect_r = min(n, int(self.Y) + 1)
        expect_s = min(max(0, n - (int(self.Y) + 1)), max(0, int(self.W) - int(self.Y)))
        if int(self.visual_kv.filled_r) != expect_r:
            raise RuntimeError(
                f"[summary-stream] RAW holds {self.visual_kv.filled_r} records after frame {t} "
                f"(n_body={n}, expected {expect_r})"
            )
        if int(self.visual_kv.filled_s) != expect_s:
            raise RuntimeError(
                f"[summary-stream] SUM holds {self.visual_kv.filled_s} summaries after frame {t} "
                f"(n_body={n}, expected {expect_s})"
            )
        n_rgb = int(rgb4.size(1))
        tot_ms = sum(ms.values())
        self._peak_ema = tot_ms if self._peak_ema is None else 0.9 * self._peak_ema + 0.1 * tot_ms
        print(
            f"[perf] abs_t={t} " + " ".join(f"{k}={v:.1f}ms" for k, v in ms.items())
            + f" total={tot_ms:.1f}ms"
            + f" | peak model+sae={1000.0 / tot_ms:.1f} lat/s = {n_rgb * 1000.0 / tot_ms:.1f} rgb fps"
            + f" (ema {n_rgb * 1000.0 / self._peak_ema:.1f})"
            + f" | x0 mean={x_mean:.3f} std={x_std:.3f} absmax={x_absmax:.2f} {self._chaser.describe()}",
            flush=True,
        )
        self.last_x0, self._abs_t, self.last_ms = xf, t + 1, dict(ms, total=tot_ms)
        return self._rgb_out(rgb4, frames=4)

    @torch.no_grad()
    def reset(self) -> None:
        """Between players (and inside start()). Zeros the one buffer in place; re-enables recency; new SinkChaser."""
        kv = self.visual_kv
        kv.reset()
        kv.enable_recency(self.spec.recency)
        self.recency = kv.recency
        self._chaser = self._new_chaser()
        if kv.filled_r or kv.filled_s or kv.n_committed or kv.z0_written:
            raise RuntimeError("[summary-stream] reset left cache state behind")
        self._rng = None
        self._abs_t, self._n_body, self._sae_cache, self.last_x0 = None, 0, None, None
        self._hist.clear()
        self._z0_gauss = None
        self._since_reprefill = 0
        self.n_reprefills = 0
        self.last_reprefill = None
        self._fps_logged = "unset"
        self._peak_ema = None
        self._mode_log.clear()
        print("[kv-boundary] reset visual=zeroed sae=dropped recency+chaser=rebuilt", flush=True)

    @torch.no_grad()
    def retarget(self, spec: StageSpec) -> None:
        """Same generator / buffer geometry; new ckpt's timesteps, stats, cache_mode, sink gap, recency, focus."""
        if not isinstance(spec, StageSpec):
            raise TypeError(f"[summary-stream] retarget spec must be StageSpec; got {type(spec)}")
        _check_spec_against_model(self.model, spec)
        if (spec.h, spec.w) != (self.spec.h, self.spec.w):
            raise RuntimeError(
                f"[summary-stream] retarget: hw {spec.h}x{spec.w} != session {self.spec.h}x{self.spec.w}; "
                f"release() and rebuild"
            )
        new_on = spec.recency is not None and bool(spec.recency.get("enabled", True))
        old_on = self.recency is not None
        if new_on != old_on and self._any_generator_compiled:
            raise RuntimeError(
                "[summary-stream] retarget: recency on<->off changes the compiled graph; release() and rebuild"
            )
        self.spec = spec
        self.focus_region_default = self.focus_region = _coerce_focus_region(spec.focus_region, spec, self.device)
        self._stats_dev = {
            "z0": _stats_to_device(spec.stats_z0, self.device),
            "rest": _stats_to_device(spec.stats_rest, self.device),
        }
        self._sqrt_ab, self._sqrt_om = _sched_tables(spec.samp, self.device)
        self.reset()
        print(
            f"[summary-stream] retargeted cache_mode={spec.cache_mode} {sink_gaps_describe(spec.sink_gap_max)} "
            f"{self._focus_region_desc()}",
            flush=True,
        )

    @torch.no_grad()
    def release(self) -> None:
        """Drop every device tensor this session owns so a switch can allocate a different buffer."""
        try:
            self.reset()
        except Exception:
            pass
        if getattr(self, "visual_kv", None) is not None:
            self.visual_kv.release()
            self.visual_kv = None
        self.recency = None
        self._sae_cache = None
        self._stats_dev = {}
        self._sqrt_ab = self._sqrt_om = None
        self._abs_t_dev = self._t_dev = self._fps_dev = self._action_dev = self._action_pin = None
        self._rgb_mean = self._rgb_std = None
        self._ev = {}
        print("[kv-boundary] released visual=dropped sae=dropped", flush=True)


def build_session_from_runtime(
    generator,
    sae_decoder,
    spec,
    device,
    *,
    rgb_hw,
    forward_body_compiled=None,
    decode_steady_compiled=None,
    use_compiled_sae_in_steady_state: bool = True,
) -> SummaryStreamSession:
    """Server: one call per served checkpoint (startup + every switch), on the CUDA thread."""
    return SummaryStreamSession(
        generator=generator,
        sae_decoder=sae_decoder,
        spec=spec,
        device=device,
        rgb_hw=rgb_hw,
        forward_body_compiled=forward_body_compiled,
        decode_steady_compiled=decode_steady_compiled,
        use_compiled_sae_in_steady_state=use_compiled_sae_in_steady_state,
    )