flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/guard.py

@@ -0,0 +1,421 @@
+"""Runtime seam contract: check it, fall back, and leave a receipt.
+
+A bound structure is calibrated for one execution form — a device, an
+input dtype, a width, and for the implementations that preallocate, a row
+count. Hand it something else and the honest outcomes are two: refuse, or
+run the host's own module instead. The outcome this module exists to
+prevent is the third one, where the kernel accepts a buffer it should not
+have and the host keeps going with a wrong answer.
+
+Falling back is allowed. Falling back *quietly* is not, because a seam
+that always falls back reads exactly like a seam that works: the plan
+still lists it, the receipt still counts it, and the latency it was
+supposed to buy is simply absent. So every fallback is counted, the first
+one per seam says so out loud, and a seam that never stops falling back
+takes itself out of the model rather than staying on as a lie.
+
+Three layers, one per timescale:
+
+  per call      the contract is checked; a violation runs the host module
+                and increments the seam's ledger
+  first time    that seam warns once, naming the path and the reason
+  persistent    after ``SELF_DETACH_AFTER`` consecutive fallbacks the seam
+                restores the host module permanently and says so
+
+``attach`` collects every guard it swaps in, so ``handle.report()`` is the
+one place to ask what actually ran. Tests assert on it: a probe that does
+not check its ledger has not checked that the thing it measured was on.
+
+The check runs in Python, so it costs a call in an eager host and nothing
+in a captured one — a graph replays the kernel it traced and never
+re-enters this code. That is the right way round: the eager host is the
+one whose shapes and devices can still change under it.
+"""
+
+from __future__ import annotations
+
+import threading
+import warnings
+from typing import Any, Callable, Iterable
+
+import torch
+
+GUARD_ATTR = "_frt_guard"
+SELF_DETACH_AFTER = 32
+
+#: returned by :meth:`GuardedSeam._frt_admit` when the call may proceed
+PROCEED = object()
+
+_FP8 = torch.float8_e4m3fn
+#: what a bf16-entry implementation will cast for itself. FP8 is
+#: deliberately absent: an fp8 tensor handed to a bf16 entry is a
+#: negotiation that did not happen, not an input to convert.
+CAST_OK = frozenset({torch.bfloat16, torch.float16, torch.float32})
+#: what an fp8-entry implementation requires, exactly
+FP8_ONLY = frozenset({_FP8})
+
+
+class GuardRefused(RuntimeError):
+    """A seam was called outside the form it was bound for."""
+
+
+def _concrete(device: torch.device | str) -> torch.device:
+    """Resolve a device to the indexed form a tensor will report.
+
+    ``torch.device("cuda")`` carries no index while every tensor on it
+    reports ``cuda:0``, so an index-less contract would refuse every call
+    it was meant to admit — a guard that fails closed on correct input is
+    worse than no guard, because it reverts a working seam and blames the
+    input. Normalising once at bind time keeps the per-call test an exact
+    comparison.
+    """
+    device = torch.device(device)
+    if device.type == "cuda" and device.index is None:
+        return torch.device("cuda", torch.cuda.current_device())
+    return device
+
+
+class SeamGuard:
+    """One seam's runtime contract, its ledger, and its own exit.
+
+    Created at bind time by the implementation, which knows the form it
+    was calibrated for. Given a site by :func:`~flashrt_structures.swap.attach`,
+    which knows where in the model it ended up.
+    """
+
+    __slots__ = ("dtypes", "device", "k", "rows", "row_capacity", "kind",
+                 "can_fallback", "calls", "fallbacks", "consecutive",
+                 "last_reason", "detached", "site", "mode", "notes",
+                 "thread", "_restore", "_warned", "pair")
+
+    def __init__(self, *, kind: str, dtypes: Iterable[torch.dtype] | None,
+                 device: torch.device, k: int | None, rows: int | None,
+                 row_capacity: int | None, can_fallback: bool) -> None:
+        self.kind = kind
+        #: structure-specific counters an implementation keeps about
+        #: itself, reported alongside the contract ones
+        self.notes: dict[str, Any] = {}
+        self.dtypes = None if dtypes is None else frozenset(dtypes)
+        self.device = _concrete(device)
+        self.k = k
+        self.rows = rows
+        self.row_capacity = row_capacity
+        self.can_fallback = can_fallback
+        self.calls = 0
+        self.fallbacks = 0
+        self.consecutive = 0
+        self.last_reason: str | None = None
+        self.detached = False
+        self.site: str | None = None
+        self.mode = "fallback"
+        #: the thread that first ran this seam. These implementations keep
+        #: preallocated stash, quantize and attention scratch buffers whose
+        #: whole economy is that one call consumes a result before the next
+        #: one writes it; a second thread in the same seam interleaves the
+        #: writes and the corruption is silent and shape-correct. Single
+        #: stream per attachment is the contract, and this is where it is
+        #: enforced rather than assumed.
+        self.thread: int | None = None
+        self._restore: Callable[[], None] | None = None
+        self._warned: set[str] = set()
+        #: the other half of a negotiated pair (a producer whose only
+        #: consumer this seam is, or vice versa). A pair is
+        #: all-or-nothing: one out-of-contract call falsifies the
+        #: premise for both seats, and a producer left seated alone
+        #: keeps feeding a host that cannot read its format.
+        self.pair: "SeamGuard | None" = None
+
+    # ---- site binding (attach time) --------------------------------
+
+    def bind_site(self, site: str, *, restore: Callable[[], None] | None,
+                  mode: str) -> None:
+        self.site = site
+        self._restore = restore
+        self.mode = mode
+
+    def release_site(self) -> None:
+        self.site = None
+        self._restore = None
+
+    # ---- per call --------------------------------------------------
+
+    def admit(self, x: torch.Tensor) -> str | None:
+        """Return ``None`` to proceed, else why this call cannot.
+
+        The fast path is the whole point: one dtype test, one width test,
+        one device test, and a row test only where the implementation
+        preallocated for a fixed row count.
+        """
+        self.calls += 1
+        if self.dtypes is not None and x.dtype not in self.dtypes:
+            return f"input dtype {x.dtype} (bound for {self._dtype_note()})"
+        if self.k is not None and x.shape[-1] != self.k:
+            return f"input width {x.shape[-1]} (bound for {self.k})"
+        if x.device != self.device:
+            return f"input on {x.device} (bound on {self.device})"
+        if self.rows is not None:
+            rows = x.numel() // x.shape[-1]
+            if rows != self.rows:
+                return f"{rows} row(s) (bound for {self.rows})"
+        if self.row_capacity is not None:
+            rows = x.numel() // x.shape[-1]
+            if rows > self.row_capacity:
+                return (f"{rows} row(s) (buffer capacity "
+                        f"{self.row_capacity})")
+        tid = threading.get_ident()
+        if self.thread is None:
+            self.thread = tid
+        elif tid != self.thread:
+            return ("called from a second thread; this seam's stash and "
+                    "scratch buffers are shared across its calls, so "
+                    "concurrent use would interleave writes into them")
+        self.consecutive = 0
+        return None
+
+    def _dtype_note(self) -> str:
+        if self.dtypes == FP8_ONLY:
+            return "fp8 from a producer seam"
+        return "/".join(sorted(str(d).replace("torch.", "")
+                               for d in self.dtypes or ()))
+
+    def refuse(self, reason: str) -> None:
+        """Record a fallback; speak up the first time and at the end."""
+        where = self.site or f"<unattached {self.kind}>"
+        if self.mode == "raise":
+            # before counting: in strict mode nothing fell back, the call
+            # was refused, and a ledger claiming otherwise would be one
+            # more thing that says something that did not happen
+            self.last_reason = reason
+            raise GuardRefused(f"{where}: {reason}")
+        self.fallbacks += 1
+        self.consecutive += 1
+        self.last_reason = reason
+        if reason not in self._warned:
+            self._warned.add(reason)
+            warnings.warn(
+                f"structures: {where} fell back to the host module — "
+                f"{reason}. Further fallbacks of this kind are counted in "
+                f"handle.report() and not repeated here.",
+                RuntimeWarning, stacklevel=4)
+        if self.pair is not None and not self.detached:
+            # a negotiated pair collapses on the first out-of-contract
+            # call: the runtime path has disproved the bind-time
+            # premise, and every call the pair survives past this one
+            # is a producer feeding a consumer that is no longer there
+            self._self_detach()
+            if not self.pair.detached:
+                self.pair._self_detach()
+            return
+        if (self.consecutive >= SELF_DETACH_AFTER and not self.detached):
+            self._self_detach()
+
+    def _self_detach(self) -> None:
+        """Stop pretending: put the host module back for good."""
+        self.detached = True
+        where = self.site or f"<unattached {self.kind}>"
+        if self._restore is None:
+            warnings.warn(
+                f"structures: {where} has fallen back "
+                f"{self.consecutive} times in a row and cannot restore "
+                f"itself (it is held inside another structure, not swapped "
+                f"in at a path). It is costing a check per call and buying "
+                f"nothing — detach the attachment.",
+                RuntimeWarning, stacklevel=5)
+            return
+        self._restore()
+        warnings.warn(
+            f"structures: {where} fell back {self.consecutive} times in a "
+            f"row ({self.last_reason}) — the host module has been restored "
+            f"at that path permanently. The rest of the attachment is "
+            f"unaffected.",
+            RuntimeWarning, stacklevel=5)
+
+    # ---- receipt ---------------------------------------------------
+
+    def entry(self) -> dict[str, Any]:
+        entry = {"kind": self.kind, "calls": self.calls,
+                 "fallbacks": self.fallbacks, "detached": self.detached,
+                 "last_reason": self.last_reason,
+                 "form": {"dtype": self._dtype_note(), "k": self.k,
+                          "rows": self.rows,
+                          "row_capacity": self.row_capacity,
+                          "device": str(self.device)}}
+        if self.notes:
+            entry["notes"] = dict(self.notes)
+        return entry
+
+
+class GuardedSeam:
+    """Mixin: contract check, host fallback, and lifecycle refusals.
+
+    Mixed in *before* ``torch.nn.Module`` so its ``_apply`` and
+    ``state_dict`` overrides win. An implementation arms itself at the end
+    of ``__init__`` with :meth:`_frt_arm` and opens its ``forward`` with
+    :meth:`_frt_admit`.
+    """
+
+    #: attribute holding the retained host module, if any
+    _frt_host_attr: str | None = None
+    #: whether calling that module with this ``forward``'s own arguments
+    #: reproduces the host's behaviour. False where the replacement's
+    #: boundary is not the retained module's boundary (a composed block,
+    #: an attention core called by an adapter): those refuse instead.
+    _frt_can_fallback: bool = False
+    #: class default so attribute lookup never reaches a host-forwarding
+    #: ``__getattr__`` for an implementation that has not armed yet
+    _frt_guard: SeamGuard | None = None
+
+    def _frt_arm(self, *, dtypes: Iterable[torch.dtype] | None,
+                 device: torch.device, k: int | None = None,
+                 rows: int | None = None,
+                 row_capacity: int | None = None) -> SeamGuard:
+        """Declare the form this instance was calibrated for."""
+        if rows is not None and row_capacity is not None:
+            raise ValueError("rows and row_capacity are mutually exclusive")
+        guard = SeamGuard(
+            kind=type(self).__name__, dtypes=dtypes, device=device, k=k,
+            rows=rows, row_capacity=row_capacity,
+            can_fallback=self._frt_can_fallback)
+        object.__setattr__(self, GUARD_ATTR, guard)
+        return guard
+
+    def _frt_touch(self) -> None:
+        """Count a call whose form another guard already checked.
+
+        Two cases need this. A replacement that ignores its inputs (a
+        frozen buffer on a slower cadence) has no contract to test but
+        still has to be countable. And a structure reached through a wider
+        entry than its own ``forward`` — a producer that a composed block
+        calls directly — is checked by the block's contract and would
+        otherwise report zero calls, which reads as "never ran" when it
+        ran every time.
+        """
+        guard = self._frt_guard
+        if guard is not None and not torch.compiler.is_compiling():
+            guard.calls += 1
+
+    def _frt_host(self) -> torch.nn.Module | None:
+        attr = self._frt_host_attr
+        if attr is None:
+            return None
+        # straight out of _modules: going through getattr would hit the
+        # host-forwarding __getattr__ these implementations define
+        return self.__dict__.get("_modules", {}).get(attr)
+
+    def _frt_admit(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Any:
+        """``PROCEED``, or the host's own result for this call.
+
+        Raises :class:`GuardRefused` when there is no host module to fall
+        back to, or when the attachment asked for strict mode.
+
+        Steps aside entirely while a compiler is tracing, for two reasons
+        that point the same way. It has to: the ledger's counter is a
+        Python-visible side effect in a hot forward, and dynamo either
+        breaks the graph around it or re-derives the region's boundaries
+        around it — the surrounding compiled region is worth more than the
+        count. And it may: dynamo installs its own guards on exactly the
+        dtype, shape and device this contract tests, so a traced region
+        handed a different form recompiles rather than running the wrong
+        kernel. Inside a compiled region this check is redundant; in an
+        eager host it is the only thing there.
+        """
+        guard = self._frt_guard
+        if guard is None or torch.compiler.is_compiling():
+            return PROCEED
+        reason = guard.admit(x)
+        if reason is None:
+            return PROCEED
+        host = self._frt_host()
+        if host is None or not guard.can_fallback:
+            raise GuardRefused(
+                f"{guard.site or guard.kind}: {reason}; this structure has "
+                "no equivalent host module to fall back to")
+        if any(p.is_meta or p.numel() == 0
+               for p in host.parameters()) and not \
+                getattr(host, "_frt_tickets", None):
+            raise GuardRefused(
+                f"{guard.site or guard.kind}: {reason}; the host module's "
+                "weights were consumed and cannot be restored")
+        if getattr(host, "_frt_tickets", None):
+            # the host's weights were consumed (their truth lives in the
+            # weight store); a live fallback restores them once — slower
+            # than a resident copy, never wrong — and the ledger says so
+            from .storage import restore_for_fallback
+            restore_for_fallback(host)
+            guard.notes["restored_for_fallback"] = (
+                guard.notes.get("restored_for_fallback", 0) + 1)
+        guard.refuse(reason)
+        return self._frt_host_call(host, x, *args, **kwargs)
+
+    def _frt_host_call(self, host: torch.nn.Module, x: torch.Tensor,
+                       *args: Any, **kwargs: Any) -> Any:
+        """How to reproduce this seam with the host module. Overridable."""
+        return host(x, *args, **kwargs)
+
+    # ---- lifecycle refusals ---------------------------------------
+
+    def _apply(self, *args: Any, **kwargs: Any):
+        """Refuse device/dtype migration while swapped into a model.
+
+        The packed weights and static scales of these implementations are
+        derived tensors held outside the parameter system, so a migration
+        would move the host's copy and leave the kernel's behind. Before
+        attach and after detach this is an ordinary module and migrates
+        normally.
+        """
+        guard = self._frt_guard
+        if guard is not None and guard.site is not None:
+            raise GuardRefused(
+                f"{guard.site}: cannot migrate device or dtype while "
+                "structures are attached — detach() first, migrate, then "
+                "attach again (the scales are calibrated per form and do "
+                "not survive a dtype change anyway)")
+        return super()._apply(*args, **kwargs)
+
+    def state_dict(self, *args: Any, destination: Any = None,
+                   prefix: str = "", keep_vars: bool = False):
+        """Emit the host module's own state at this path.
+
+        A checkpoint should not change shape because an optimisation is
+        attached: the packed and quantised tensors here are derived from
+        the host weights and belong to the bind, not to the model's
+        state. Delegating keeps the schema and the values identical to the
+        unattached model, so saving while attached is safe.
+        """
+        host = self._frt_host()
+        if host is None or args:
+            return super().state_dict(*args, destination=destination,
+                                      prefix=prefix, keep_vars=keep_vars)
+        return host.state_dict(destination=destination, prefix=prefix,
+                               keep_vars=keep_vars)
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata,
+                              strict, missing_keys, unexpected_keys,
+                              error_msgs):
+        """Refuse loading into a swapped model rather than half-doing it.
+
+        Loading would have to re-derive every packed weight and re-run
+        calibration; silently loading the host's copy and leaving the
+        kernel's stale is the failure this refuses. ``detach()`` restores
+        the real modules, so the honest order is detach, load, attach.
+        """
+        guard = self._frt_guard
+        if guard is not None and guard.site is not None:
+            error_msgs.append(
+                f"{guard.site}: cannot load_state_dict while structures "
+                "are attached — detach(), load, then attach again so the "
+                "packed weights are rebuilt from the loaded ones")
+            return
+        return super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys,
+            unexpected_keys, error_msgs)
+
+
+def collect(module: torch.nn.Module) -> list[tuple[str, SeamGuard]]:
+    """Every guard at or under ``module``, by relative module name."""
+    found = []
+    for name, child in module.named_modules():
+        guard = child.__dict__.get(GUARD_ATTR)
+        if isinstance(guard, SeamGuard):
+            found.append((name, guard))
+    return found

flashrt_structures/handle.py

@@ -0,0 +1,189 @@
+"""Explicit front door: ``structures.get(name)`` — pull one structure,
+bind it to one host module, plug it in yourself.
+
+This is the single-site, fully visible counterpart of ``attach``. The
+consumption mirrors ``kernels.get_kernel``:
+
+    ffn = structures.get("decoder_ffn")
+    new_mlp = ffn.bind(model.model.layers[3].mlp, calibration=[x1, x2])
+    model.model.layers[3].mlp = new_mlp        # your swap, your call
+
+``bind`` extracts the weights from the module you hand it, calibrates on
+the samples you provide, builds the fused replacement, and self-checks
+it against the original module on those same samples. A replacement
+that misses the parity gate raises ``GateRefused`` instead of returning
+— the explicit path fails loudly, it never hands back a bad part.
+``attach`` is this same operation, batched over every discovered site
+with the net-win A/B added.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import torch
+
+from .discover import _VISION_PROJ, activation_for
+from .gates import parity_metrics
+from flash_rt.catalog.registry import StructureSpec, load
+
+
+class GateRefused(RuntimeError):
+    """The bound replacement did not meet the parity gate."""
+
+
+class StructureHandle:
+    """One pulled structure, ready to bind to host modules."""
+
+    def __init__(self, spec: StructureSpec):
+        self.spec = spec
+        self.name = spec.name
+
+    def __repr__(self) -> str:
+        return f"StructureHandle({self.name!r}, v{self.spec.version})"
+
+    def bind(
+        self,
+        module: torch.nn.Module,
+        calibration: torch.Tensor | Sequence[torch.Tensor],
+        *,
+        variant: dict[str, str] | None = None,
+        norm: torch.nn.Module | None = None,
+        residual: torch.Tensor | Sequence[torch.Tensor] | None = None,
+        gate_cos: float = 0.999,
+        percentile: float = 99.9,
+    ) -> torch.nn.Module:
+        """Build a gated drop-in replacement for ``module``.
+
+        ``calibration``: one or more real inputs of ``module`` (the
+        normed hidden states it actually sees). ``residual``: the
+        matching pre-norm hidden states — the structure gate is declared
+        at the boundary *including* the residual add, so provide these
+        to gate at the declared boundary; without them the self-check
+        runs on the bare seam output, which is strictly harsher.
+        ``norm``: optional sibling norm module, recorded for
+        full-structure use; the MLP seam itself leaves the norm in the
+        host. ``gate_cos=0`` skips the self-check (not recommended).
+        """
+        samples = ([calibration] if torch.is_tensor(calibration)
+                   else list(calibration))
+        if not samples:
+            raise ValueError("calibration must contain at least one sample")
+        residuals = (None if residual is None else
+                     [residual] if torch.is_tensor(residual)
+                     else list(residual))
+        if residuals is not None and len(residuals) != len(samples):
+            raise ValueError("residual must align 1:1 with calibration")
+
+        if self.name == "decoder_ffn":
+            from .impls.decoder_ffn import fp8_static as impl
+            gate = module.gate_proj
+            device = gate.weight.device
+            dims = {"D": gate.in_features, "F": gate.out_features}
+            var = {"activation": activation_for(module, "silu"),
+                   "norm_weight_mode": "direct"}
+            weights = {
+                "w_norm": (norm.weight.detach() if norm is not None
+                           and getattr(norm, "weight", None) is not None
+                           else torch.ones(dims["D"], device=device)),
+                "w_gate": gate.weight.detach().t().contiguous(),
+                "w_up": module.up_proj.weight.detach().t().contiguous(),
+                "w_down": module.down_proj.weight.detach().t().contiguous(),
+            }
+        elif self.name == "vision_ffn":
+            from .impls.vision_ffn import fp8_static as impl
+            fc_attrs = next(
+                (pair for pair in _VISION_PROJ
+                 if all(isinstance(getattr(module, a, None), torch.nn.Linear)
+                        for a in pair)), None)
+            if fc_attrs is None:
+                raise ValueError(
+                    "module has no fc1/fc2 (or linear_fc1/linear_fc2) pair")
+            fc1 = getattr(module, fc_attrs[0])
+            fc2 = getattr(module, fc_attrs[1])
+            device = fc1.weight.device
+            dims = {"D": fc1.in_features, "F": fc1.out_features}
+            var = {"activation": activation_for(module, "gelu")}
+            weights = {
+                "w_norm": (norm.weight.detach() if norm is not None
+                           else torch.ones(dims["D"], device=device)),
+                "b_norm": (norm.bias.detach() if norm is not None
+                           else torch.zeros(dims["D"], device=device)),
+                "w_fc1": fc1.weight.detach(), "b_fc1": fc1.bias.detach(),
+                "w_fc2": fc2.weight.detach(), "b_fc2": fc2.bias.detach(),
+            }
+        else:
+            raise ValueError(
+                f"structure {self.name!r} has no module-seam bind; "
+                "stage pipelines go through the capture door")
+        var.update(variant or {})
+
+        # the explicit door owns real samples, so it measures the seam's
+        # points by running the host module under the same hooks the
+        # distribution layer uses — one path for what a point is and where
+        # it lives, whichever door asked
+        from .points import Collector, Point
+
+        points = [Point("x_after_norm", "", "input")]
+        second = ("down_proj" if self.name == "decoder_ffn"
+                  else (fc_attrs[1] if self.name == "vision_ffn" else None))
+        second_name = ("act_after_mul" if self.name == "decoder_ffn"
+                       else "hidden_after_act")
+        if second is not None:
+            points.append(Point(second_name, second, "input"))
+        collector = Collector(points=points)
+        handles = collector.hooks(
+            lambda path: module if not path else getattr(module, path))
+        try:
+            with torch.no_grad():
+                for sample in samples:
+                    module(sample.to(device))
+                    collector.end_sample()
+        finally:
+            for handle in handles:
+                handle.remove()
+        collector.reduce(percentile)
+        kwargs = {} if self.name == "vision_ffn" else {"variant": var}
+        bound = impl.bind_mlp_seam(
+            weights,
+            input_scale=collector.scale("", "x_after_norm"),
+            hidden_scale=collector.scale(second or "", second_name),
+            original=module, **kwargs)
+
+        worst = 1.0
+        with torch.no_grad():
+            for k, sample in enumerate(samples):
+                x = sample.to(device)
+                got, want = bound(x), module(x)
+                if residuals is not None:
+                    res = residuals[k].to(device)
+                    got, want = res + got, res + want
+                worst = min(worst, parity_metrics(got, want)["cosine"])
+        boundary = ("structure (incl. residual)" if residuals is not None
+                    else "bare seam")
+        if gate_cos and worst < gate_cos:
+            raise GateRefused(
+                f"{self.name} replacement worst cos {worst:.6f} < "
+                f"{gate_cos} at {boundary} boundary on {len(samples)} "
+                "calibration sample(s) — not handing back a part that "
+                "fails its own gate"
+                + ("" if residuals is not None else
+                   "; pass residual= to gate at the declared boundary"))
+        bound.certification = {
+            "structure": self.name, "version": self.spec.version,
+            "dims": dims, "variant": var, "worst_cos": round(worst, 7),
+            "gate_boundary": boundary, "samples": len(samples),
+            "m_profile": sorted({
+                int(s.reshape(-1, s.shape[-1]).shape[0]) for s in samples}),
+        }
+        return bound
+
+
+def get(name: str) -> StructureHandle:
+    """Pull one structure from the catalog (get_kernel-style)."""
+    spec = load(name)
+    if spec.kind != "region":
+        raise ValueError(
+            f"{name!r} is a {spec.kind} structure — module-level bind "
+            "does not apply; use the capture/provider route")
+    return StructureHandle(spec)