flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/stages.py

@@ -0,0 +1,298 @@
+"""Capture door: ``structures.capture(fn, ...)`` — graph a hot stage.
+
+The schedule-structure counterpart of the region doors. ``fn`` is the
+hot stage of a cond_iter pipeline (a denoise loop, a decode step chain);
+``capture`` warms it, records it into a CUDA graph on a side stream, and
+returns a runtime whose ``replay()`` re-executes the whole stage as one
+launch. Declared ``windows`` are the tensors you may rewrite between
+replays (noise, observations, condition buffers) — replay reads their
+current contents in place, which is what makes the graph a reusable
+stage rather than a frozen trace.
+
+Gates are built in: ``reference`` (an eager thunk producing the same
+output under the same window contents) certifies parity, and the replay
+is timed against eager ``fn``; a capture that is not both accurate and
+net-faster raises ``CaptureRefused`` — the host keeps its eager path,
+never a silently wrong or slower graph.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Mapping
+
+import torch
+
+from .gates import parity_metrics
+
+
+class CaptureRefused(RuntimeError):
+    """The captured stage failed its parity or net-win gate."""
+
+
+def _first_tensor(value: Any) -> torch.Tensor | None:
+    """Pick one tensor deterministically for schedule-equivalence preflight."""
+    if torch.is_tensor(value):
+        return value
+    logits = getattr(value, "logits", None)
+    if torch.is_tensor(logits):
+        return logits
+    if isinstance(value, Mapping):
+        for key in sorted(value):
+            found = _first_tensor(value[key])
+            if found is not None:
+                return found
+        return None
+    if isinstance(value, (tuple, list)):
+        for item in value:
+            found = _first_tensor(item)
+            if found is not None:
+                return found
+    return None
+
+
+def _time_ms(fn: Callable[[], Any], warmup: int = 5, iters: int = 20) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start, end = torch.cuda.Event(True), torch.cuda.Event(True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+@dataclass
+class CapturedStage:
+    """A replayable hot stage with declared swap windows."""
+
+    graph: torch.cuda.CUDAGraph
+    stream: torch.cuda.Stream
+    output: Any
+    windows: Mapping[str, torch.Tensor]
+    certification: dict[str, Any] = field(default_factory=dict)
+    #: family shape-lowerings applied to the host before capture. Replay
+    #: does not need them (a graph replays kernels, not Python), but the
+    #: host object stays pinned for any eager use until they are undone.
+    lowerings: tuple = ()
+
+    def restore_host(self) -> None:
+        """Undo every family lowering; the host runs eager as loaded."""
+        for lowering in reversed(self.lowerings):
+            lowering.undo()
+        self.lowerings = ()
+
+    def replay(self, sync: bool = True) -> Any:
+        with torch.cuda.stream(self.stream):
+            self.graph.replay()
+        if sync:
+            torch.cuda.synchronize()
+        return self.output
+
+    def write(self, name: str, value: torch.Tensor) -> None:
+        """Rewrite one declared window in place.
+
+        The window's form is the graph's contract: ``copy_`` would
+        silently broadcast a wrong shape or cast a wrong dtype, and the
+        replay would then run baked kernels over coerced data with no
+        guard anywhere to notice — the captured tier's ledger blindness
+        is by design, so the door has to check.
+        """
+        window = self.windows[name]
+        if (value.shape != window.shape or value.dtype != window.dtype
+                or value.device != window.device):
+            raise ValueError(
+                f"captured stage: window {name!r} expects "
+                f"{tuple(window.shape)}/{window.dtype}/{window.device}, "
+                f"got {tuple(value.shape)}/{value.dtype}/{value.device}"
+                " — a replay over a coerced write is silently wrong")
+        window.copy_(value)
+
+    def export(self, *, ports, regions=(), stages=None, roles=None,
+               identity=None, manifest_extra=None):
+        """One call from this captured stage to ``frt_model_runtime_v1``.
+
+        The declared windows become the contract's boundary windows
+        (name -> device pointer + bytes); ``ports`` reference them by
+        name exactly as in
+        :func:`flash_rt.runtime.provider.export_captured_runtime`.
+        This is the absorption edge: a stage captured out of any torch
+        host becomes a runtime the FlashRT serving mechanisms (Nexus
+        tick, capsule snapshot/restore) consume like a whitebox one.
+        """
+        from flash_rt.runtime.provider import export_captured_runtime
+
+        window_map = {
+            name: (t.data_ptr(), t.numel() * t.element_size())
+            for name, t in self.windows.items()}
+        return export_captured_runtime(
+            stream_handle=self.stream.cuda_stream,
+            graphs=[("infer", self.graph.raw_cuda_graph_exec())],
+            windows=window_map,
+            ports=ports, regions=regions,
+            stages=tuple(stages) if stages else ("infer",),
+            roles=roles, identity=identity,
+            manifest_extra=manifest_extra)
+
+
+def capture(
+    fn: Callable[[], Any],
+    *,
+    model: torch.nn.Module | None = None,
+    windows: Mapping[str, torch.Tensor] | None = None,
+    reference: Callable[[], torch.Tensor] | None = None,
+    output_of: Callable[[Any], torch.Tensor] | None = None,
+    warmup: int = 3,
+    gate_cos: float = 0.999,
+    min_speedup: float = 1.02,
+    verbose: bool = True,
+) -> CapturedStage:
+    """Capture ``fn`` into a replayable, gated stage.
+
+    ``fn`` must be graph-safe: fixed shapes, no host-side branching on
+    tensor values, all varying inputs read from the declared ``windows``
+    buffers. When ``model`` is supplied, a registered fixed-iteration
+    schedule adapter may normalize a recognized graph-unsafe host loop; the
+    original and normalized callables must pass an exact preflight before
+    capture. ``reference`` runs the host's own eager path under the same
+    window contents; parity is judged between its output and the
+    replayed stage output. Pass ``gate_cos=0`` / ``min_speedup=0`` to
+    skip a gate explicitly (recorded in the certification).
+    """
+
+    def say(msg: str) -> None:
+        if verbose:
+            print(f"[structures] {msg}", flush=True)
+
+    from .impls.fixed_iter import (
+        FixedIterationRefused,
+        normalize_fixed_iteration,
+    )
+
+    schedule = None
+    try:
+        schedule = normalize_fixed_iteration(fn, model)
+    except FixedIterationRefused as exc:
+        raise CaptureRefused(str(exc)) from exc
+    if schedule is not None:
+        want_schedule = _first_tensor(schedule.reference_output)
+        with torch.no_grad():
+            got_schedule = _first_tensor(schedule.forward())
+        if want_schedule is None or got_schedule is None:
+            raise CaptureRefused(
+                f"{schedule.family}: schedule parity needs one tensor output")
+        metrics = parity_metrics(
+            got_schedule.detach().float().cpu(),
+            want_schedule.detach().float().cpu(),
+        )
+        exact = torch.equal(got_schedule, want_schedule)
+        if schedule.exact and not exact:
+            raise CaptureRefused(
+                f"{schedule.family}: fixed-iteration lowering is not exact "
+                f"(cos={metrics['cosine']:.7f}, "
+                f"max_abs={metrics['max_abs']:.7g})")
+        fn = schedule.forward
+        if schedule.compile_before_capture:
+            fn = torch.compile(fn)
+        say(f"schedule normalized: {schedule.family}, "
+            f"{schedule.steps} fixed step(s), exact={exact}")
+
+    lowerings: list = []
+    if model is not None:
+        from .impls.graph_lowering import lower_for_capture
+        lowerings = lower_for_capture(model, fn)
+        for lowering in lowerings:
+            say(f"graph lowering applied: {lowering.family} "
+                f"[{', '.join(lowering.pins)}]")
+
+    declared_windows = dict(schedule.windows) if schedule is not None else {}
+    for name, tensor in dict(windows or {}).items():
+        if name in declared_windows \
+                and declared_windows[name] is not tensor:
+            raise CaptureRefused(
+                f"window {name!r} conflicts with the tensor discovered by "
+                "the fixed-iteration schedule")
+        declared_windows[name] = tensor
+    windows = declared_windows
+    stream = torch.cuda.Stream()
+    try:
+        with torch.no_grad(), torch.cuda.stream(stream):
+            for _ in range(max(1, warmup)):
+                fn()
+            torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        # same grad mode as the warmups: a mode switch here changes the
+        # dynamo guards and a compiled fn recompiles inside the capture,
+        # which runs CUDA work no capture may record
+        with torch.no_grad(), torch.cuda.graph(graph, stream=stream):
+            output = fn()
+        torch.cuda.synchronize()
+    except Exception:
+        for lowering in reversed(lowerings):
+            lowering.undo()
+        raise
+    say(f"stage captured ({len(windows)} window(s))")
+
+    stage = CapturedStage(graph=graph, stream=stream, output=output,
+                          windows=windows, lowerings=tuple(lowerings))
+    pick = output_of or (lambda out: out)
+
+    cert: dict[str, Any] = {"windows": sorted(windows),
+                            "gate_cos": gate_cos,
+                            "min_speedup": min_speedup}
+    if lowerings:
+        cert["graph_lowering"] = [
+            {"family": low.family, "pins": list(low.pins),
+             **dict(low.details)} for low in lowerings]
+    if schedule is not None:
+        cert["schedule"] = {
+            "family": schedule.family,
+            "steps": schedule.steps,
+            "exact": exact,
+            "parity_cos": round(metrics["cosine"], 7),
+            "max_abs": metrics["max_abs"],
+            "compiled": schedule.compile_before_capture,
+            **dict(schedule.details),
+        }
+    if reference is not None and gate_cos:
+        with torch.no_grad():
+            want = reference().detach().float().cpu()
+        got = pick(stage.replay()).detach().float().cpu()
+        cos = parity_metrics(got, want)["cosine"]
+        cert["parity_cos"] = round(cos, 7)
+        if cos < gate_cos:
+            raise CaptureRefused(
+                f"captured stage parity cos {cos:.6f} < {gate_cos} vs "
+                "eager reference — check that every varying input is a "
+                "declared window (in-graph RNG never matches eager)")
+        say(f"parity vs eager: cos={cos:.6f}")
+    if min_speedup:
+        eager_ms = _time_ms(lambda: fn())
+        # replays run on the side stream: the timing events must be
+        # recorded on that same stream or they only measure enqueue time
+        for _ in range(5):
+            stage.replay(sync=False)
+        torch.cuda.synchronize()
+        start, end = torch.cuda.Event(True), torch.cuda.Event(True)
+        with torch.cuda.stream(stream):
+            start.record()
+        for _ in range(20):
+            stage.replay(sync=False)
+        with torch.cuda.stream(stream):
+            end.record()
+        torch.cuda.synchronize()
+        replay_ms = start.elapsed_time(end) / 20
+        cert["eager_ms"] = round(eager_ms, 3)
+        cert["replay_ms"] = round(replay_ms, 3)
+        cert["speedup"] = round(eager_ms / replay_ms, 4)
+        if eager_ms / replay_ms < min_speedup:
+            raise CaptureRefused(
+                f"captured stage is not a net win "
+                f"({eager_ms / replay_ms:.3f}x vs eager, margin "
+                f"{min_speedup}) — host keeps its eager path")
+        say(f"net win: {eager_ms:.2f} -> {replay_ms:.2f} ms "
+            f"({eager_ms / replay_ms:.3f}x)")
+    stage.certification = cert
+    return stage

flashrt_structures/storage.py

@@ -0,0 +1,255 @@
+"""Weight residency: the fallback store lives off the GPU.
+
+The attachment used to keep every replaced module's parameters resident
+in device memory so that per-call fallback and ``detach`` were instant.
+That is a double weight bill paid exactly when memory is scarcest — at
+bind time on a large host — and the only thing it bought was speed of a
+path that production (captured, fixed-shape) never takes.
+
+This module removes the resident tier. When an attachment consumes its
+originals, each parameter's truth moves to one of two stores and the
+device storage is freed:
+
+``DISK``
+    The parameter verifiably came from a checkpoint file: its
+    safetensors shard and key are recorded, and a sampled-block
+    comparison against the live tensor proved the mapping before
+    anything was released. Restore reloads from the file and re-applies
+    the load transform (a dtype cast). Zero extra copies held.
+
+``HOST_RAM``
+    No verifiable provenance (the host mutated the weight after
+    loading, or the mapping could not be proven). The tensor is spilled
+    to pinned CPU memory; restore is one host-to-device copy. Costs
+    host RAM, never device memory.
+
+Correctness never depends on the mapping heuristics: a provenance miss
+costs RAM, not accuracy. Tied parameters (shared storage) are stashed
+once and restored to every holder.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+_SAMPLE_BYTES = 65536
+
+
+def _sample_equal(a: torch.Tensor, b: torch.Tensor) -> bool:
+    """Compare shape, dtype and the first/last sample blocks."""
+    if a.shape != b.shape or a.dtype != b.dtype:
+        return False
+    fa = a.reshape(-1)
+    fb = b.reshape(-1)
+    n = min(fa.numel(), _SAMPLE_BYTES // max(1, fa.element_size()))
+    if n == 0:
+        return True
+    return (torch.equal(fa[:n].cpu(), fb[:n].cpu())
+            and torch.equal(fa[-n:].cpu(), fb[-n:].cpu()))
+
+
+@dataclass
+class _Ticket:
+    tier: str                     # "disk" | "ram"
+    shape: torch.Size
+    dtype: torch.dtype
+    device: torch.device
+    shard: str | None = None      # disk tier
+    key: str | None = None
+    spill: torch.Tensor | None = None   # ram tier
+    storage_key: int = 0          # data_ptr at stash time (tied weights)
+
+
+@dataclass
+class WeightStore:
+    """Tiered off-device store for consumed host weights."""
+
+    checkpoint: str | None = None
+    _index: dict[str, str] | None = field(default=None, repr=False)
+    _by_storage: dict[int, _Ticket] = field(default_factory=dict,
+                                            repr=False)
+    _stashed: list = field(default_factory=list, repr=False)
+    stats: dict[str, Any] = field(default_factory=lambda: {
+        "disk": 0, "ram": 0, "freed_bytes": 0, "restored": 0})
+
+    # ---- provenance -------------------------------------------------
+
+    def _load_index(self) -> dict[str, str]:
+        if self._index is not None:
+            return self._index
+        index: dict[str, str] = {}
+        if self.checkpoint and os.path.isdir(self.checkpoint):
+            idx_path = os.path.join(
+                self.checkpoint, "model.safetensors.index.json")
+            if os.path.exists(idx_path):
+                with open(idx_path) as f:
+                    weight_map = json.load(f).get("weight_map", {})
+                index = {k: os.path.join(self.checkpoint, v)
+                         for k, v in weight_map.items()}
+            else:
+                for shard in glob.glob(
+                        os.path.join(self.checkpoint, "*.safetensors")):
+                    from safetensors import safe_open
+                    with safe_open(shard, framework="pt") as f:
+                        for k in f.keys():
+                            index[k] = shard
+        self._index = index
+        return index
+
+    def _disk_source(self, name: str,
+                     tensor: torch.Tensor) -> tuple[str, str] | None:
+        """A verified (shard, key) for this parameter, or None."""
+        index = self._load_index()
+        if not index:
+            return None
+        candidates = [name]
+        # hosts commonly hang the checkpoint tree under one wrapper
+        # attribute; try progressively stripped prefixes
+        parts = name.split(".")
+        for i in range(1, min(4, len(parts))):
+            candidates.append(".".join(parts[i:]))
+        from safetensors import safe_open
+        for key in candidates:
+            shard = index.get(key)
+            if shard is None:
+                continue
+            with safe_open(shard, framework="pt") as f:
+                disk = f.get_tensor(key)
+            if _sample_equal(disk.to(tensor.dtype), tensor):
+                return shard, key
+        return None
+
+    # ---- stash / restore -------------------------------------------
+
+    @torch.no_grad()
+    def stash_module(self, name: str, module: torch.nn.Module) -> int:
+        """Move every parameter's truth off the device and free it.
+
+        Returns the number of device bytes freed. Safe to call twice —
+        already-emptied parameters are skipped. The module object stays
+        whole (shapes, dtypes and attributes remain introspectable);
+        only the storage leaves.
+        """
+        freed = 0
+        for mpath, sub in module.named_modules():
+            for leaf, par in list(sub._parameters.items()):
+                if par is None or par.is_meta or par.numel() == 0:
+                    continue
+                pname = f"{mpath}.{leaf}" if mpath else leaf
+                skey = par.data.data_ptr()
+                ticket = self._by_storage.get(skey)
+                if ticket is None:
+                    full = f"{name}.{pname}" if name else pname
+                    source = self._disk_source(full, par.data)
+                    if source is not None:
+                        ticket = _Ticket(
+                            tier="disk", shape=par.shape, dtype=par.dtype,
+                            device=par.device, shard=source[0],
+                            key=source[1], storage_key=skey)
+                        self.stats["disk"] += 1
+                    else:
+                        spill = torch.empty(
+                            par.shape, dtype=par.dtype, device="cpu",
+                            pin_memory=torch.cuda.is_available())
+                        spill.copy_(par.data)
+                        ticket = _Ticket(
+                            tier="ram", shape=par.shape, dtype=par.dtype,
+                            device=par.device, spill=spill, storage_key=skey)
+                        self.stats["ram"] += 1
+                    self._by_storage[skey] = ticket
+                tickets = getattr(module, "_frt_tickets", None)
+                if tickets is None:
+                    tickets = {}
+                    module._frt_tickets = tickets
+                    module._frt_store = self
+                    self._stashed.append(module)
+                tickets[pname] = ticket
+                freed += par.numel() * par.element_size()
+                # release to a meta parameter of the SAME shape: the
+                # storage is gone, but shape probes (an attention family
+                # deriving head counts from a weight) still see the truth
+                # — compute on it fails loudly and per-site, instead of
+                # collapsing a whole discovery stage on a 1-D empty
+                sub._parameters[leaf] = torch.nn.Parameter(
+                    torch.empty(par.shape, dtype=par.dtype,
+                                device="meta"), requires_grad=False)
+        self.stats["freed_bytes"] += freed
+        return freed
+
+    @torch.no_grad()
+    def restore_module(self, module: torch.nn.Module) -> bool:
+        """Put every stashed parameter back on its device. Idempotent."""
+        tickets = getattr(module, "_frt_tickets", None)
+        if not tickets:
+            return False
+        subs = dict(module.named_modules())
+        for pname, ticket in tickets.items():
+            mpath, _, leaf = pname.rpartition(".")
+            sub = subs.get(mpath)
+            par = None if sub is None else sub._parameters.get(leaf)
+            if par is None or not par.is_meta:
+                continue
+            if ticket.tier == "disk":
+                from safetensors import safe_open
+                with safe_open(ticket.shard, framework="pt") as f:
+                    data = f.get_tensor(ticket.key)
+                data = data.to(ticket.device, ticket.dtype)
+            else:
+                data = ticket.spill.to(ticket.device, non_blocking=False)
+            if data.shape != ticket.shape:
+                raise RuntimeError(
+                    f"weight store: {pname} restored to {tuple(data.shape)}"
+                    f", expected {tuple(ticket.shape)}")
+            sub._parameters[leaf] = torch.nn.Parameter(
+                data, requires_grad=False)
+            self.stats["restored"] += 1
+        del module._frt_tickets
+        del module._frt_store
+        return True
+
+    def restore_all(self) -> int:
+        """Abort path: put every stashed module's weights back.
+
+        For a bind or attach that dies midway — the model returns to
+        runnable, whatever the failure was.
+        """
+        return _restore_all_of(self)
+
+    def drop_module(self, module: torch.nn.Module) -> None:
+        """Forget a module's tickets — its consumption becomes final."""
+        if hasattr(module, "_frt_tickets"):
+            del module._frt_tickets
+        if hasattr(module, "_frt_store"):
+            del module._frt_store
+
+
+def _restore_all_of(store: "WeightStore") -> int:
+    count = 0
+    for module in list(store._stashed):
+        if store.restore_module(module):
+            count += 1
+    store._stashed.clear()
+    return count
+
+
+def restore_for_fallback(host: torch.nn.Module) -> bool:
+    """Bring a consumed host module back for a live fallback.
+
+    Called from the guard path when a seam outside its contract needs
+    the host and the host's weights were consumed. Returns True when a
+    restore happened. Consumption made final (tickets dropped) leaves
+    nothing to restore — the caller refuses instead.
+    """
+    store = getattr(host, "_frt_store", None)
+    if store is None:
+        return False
+    return store.restore_module(host)
+
+
+__all__ = ["WeightStore", "restore_for_fallback"]