flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/attention_core/two_way_fa2.py

@@ -0,0 +1,220 @@
+"""FA2 implementation for factored two-way attention.
+
+This is the common MoT form used by multimodal diffusion transformers:
+the causal/understanding branch attends to itself causally, while the
+full/generation branch attends to the joint causal + full sequence. The
+host has already projected, normalized, and rotated Q/K/V; this structure
+owns only the two attention calls and the factored output boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardRefused, GuardedSeam
+
+
+@dataclass
+class _AttentionScratch:
+    out: torch.Tensor
+    lse: torch.Tensor
+    workspace: Any
+
+
+def _scratch(kernel, q: torch.Tensor, k: torch.Tensor) -> _AttentionScratch:
+    out, lse = kernel.allocate_outputs(q)
+    return _AttentionScratch(
+        out=out,
+        lse=lse,
+        workspace=kernel.allocate_workspace(q, k),
+    )
+
+
+class FactoredTwoWayAttention(GuardedSeam, torch.nn.Module):
+    """Allocation-free single-sample two-way GQA attention processor."""
+
+    _frt_can_fallback = False
+
+    def __init__(self, capture: dict[str, Any]) -> None:
+        super().__init__()
+        query = capture["query"]
+        key = capture["key"]
+        value = capture["value"]
+        causal_q = query["causal_seq"]
+        full_q = query["full_only_seq"]
+        causal_k = key["causal_seq"]
+        full_k = key["full_only_seq"]
+        causal_v = value["causal_seq"]
+        full_v = value["full_only_seq"]
+
+        if query["sample_offsets"].numel() != 2:
+            raise ValueError(
+                "attention_core two_way: only one-sample factored packs "
+                "are qualified")
+        if causal_q.ndim != 3 or full_q.ndim != 3:
+            raise ValueError(
+                "attention_core two_way: Q must have shape [tokens, heads, dim]")
+        if causal_k.shape != causal_v.shape or full_k.shape != full_v.shape:
+            raise ValueError(
+                "attention_core two_way: K and V shapes differ")
+        if causal_q.shape[-1] != causal_k.shape[-1]:
+            raise ValueError(
+                "attention_core two_way: Q and KV head dimensions differ")
+        if causal_q.dtype != torch.bfloat16:
+            raise ValueError(
+                "attention_core two_way: current Hub FA2 path requires BF16")
+
+        self.causal_shape = tuple(causal_q.shape)
+        self.full_q_shape = tuple(full_q.shape)
+        self.causal_kv_shape = tuple(causal_k.shape)
+        self.full_kv_shape = tuple(full_k.shape)
+        self.total_tokens = int(
+            query["_causal_indices"].numel()
+            + query["_full_indices"].numel())
+        self.scale = float(causal_q.shape[-1] ** -0.5)
+        self._kernel = hub_kernel(
+            "flashrt/fa2-seqused-runtime",
+            ">=1",
+        )
+
+        device = causal_q.device
+        dtype = causal_q.dtype
+        kv_heads = causal_k.shape[1]
+        head_dim = causal_k.shape[2]
+        self.register_buffer(
+            "joint_k",
+            torch.empty(
+                1,
+                self.total_tokens,
+                kv_heads,
+                head_dim,
+                device=device,
+                dtype=dtype,
+            ),
+            persistent=False,
+        )
+        self.register_buffer(
+            "joint_v",
+            torch.empty_like(self.joint_k),
+            persistent=False,
+        )
+        self.register_buffer(
+            "causal_indices",
+            query["_causal_indices"].long().detach().clone(),
+            persistent=False,
+        )
+        self.register_buffer(
+            "full_indices",
+            query["_full_indices"].long().detach().clone(),
+            persistent=False,
+        )
+
+        causal_q4 = causal_q.unsqueeze(0)
+        causal_k4 = causal_k.unsqueeze(0)
+        full_q4 = full_q.unsqueeze(0)
+        self._causal = _scratch(
+            self._kernel,
+            causal_q4,
+            causal_k4,
+        )
+        self._full = _scratch(
+            self._kernel,
+            full_q4,
+            self.joint_k,
+        )
+        self._frt_arm(
+            dtypes=(dtype,),
+            device=device,
+            k=int(head_dim),
+            rows=int(causal_q.shape[0] * causal_q.shape[1]),
+        )
+
+    def _validate(self, query, key, value) -> None:
+        shapes = (
+            tuple(query["causal_seq"].shape),
+            tuple(query["full_only_seq"].shape),
+            tuple(key["causal_seq"].shape),
+            tuple(key["full_only_seq"].shape),
+            tuple(value["causal_seq"].shape),
+            tuple(value["full_only_seq"].shape),
+        )
+        expected = (
+            self.causal_shape,
+            self.full_q_shape,
+            self.causal_kv_shape,
+            self.full_kv_shape,
+            self.causal_kv_shape,
+            self.full_kv_shape,
+        )
+        if shapes != expected:
+            raise GuardRefused(
+                f"attention_core two_way: shapes {shapes} "
+                f"(bound for {expected})")
+
+    def forward(self, query, key, value):
+        admitted = self._frt_admit(query["causal_seq"])
+        if admitted is not PROCEED:
+            return admitted
+        self._validate(query, key, value)
+
+        causal_q = query["causal_seq"].unsqueeze(0)
+        causal_k = key["causal_seq"].unsqueeze(0)
+        causal_v = value["causal_seq"].unsqueeze(0)
+        cs = self._causal
+        causal_out = self._kernel.forward_static(
+            causal_q,
+            causal_k,
+            causal_v,
+            out=cs.out,
+            softmax_lse=cs.lse,
+            workspace=cs.workspace,
+            softmax_scale=self.scale,
+            causal=True,
+        )
+
+        self.joint_k[0].index_copy_(
+            0,
+            self.causal_indices,
+            key["causal_seq"],
+        )
+        self.joint_k[0].index_copy_(
+            0,
+            self.full_indices,
+            key["full_only_seq"],
+        )
+        self.joint_v[0].index_copy_(
+            0,
+            self.causal_indices,
+            value["causal_seq"],
+        )
+        self.joint_v[0].index_copy_(
+            0,
+            self.full_indices,
+            value["full_only_seq"],
+        )
+        fs = self._full
+        full_out = self._kernel.forward_static(
+            query["full_only_seq"].unsqueeze(0),
+            self.joint_k,
+            self.joint_v,
+            out=fs.out,
+            softmax_lse=fs.lse,
+            workspace=fs.workspace,
+            softmax_scale=self.scale,
+            causal=False,
+        )
+
+        out = dict(query)
+        out["causal_seq"] = causal_out.squeeze(0).flatten(-2, -1)
+        out["full_only_seq"] = full_out.squeeze(0).flatten(-2, -1)
+        return out
+
+
+def bind_two_way_attention(capture: dict[str, Any]
+                           ) -> FactoredTwoWayAttention:
+    """Bind a factored two-way attention processor from one real call."""
+    return FactoredTwoWayAttention(capture)

flashrt_structures/impls/cadence_static/__init__.py

@@ -0,0 +1,18 @@
+from .buffers import StaticOutput, bind_cadence_static
+from .cross_attention import (
+    CrossKvCandidate,
+    bind_cross_attention_kv,
+    capture_cross_attention_kv,
+    discover_cross_attention_kv,
+    refresh_cross_attention_kv,
+)
+
+__all__ = [
+    "CrossKvCandidate",
+    "StaticOutput",
+    "bind_cadence_static",
+    "bind_cross_attention_kv",
+    "capture_cross_attention_kv",
+    "discover_cross_attention_kv",
+    "refresh_cross_attention_kv",
+]

flashrt_structures/impls/cadence_static/buffers.py

@@ -0,0 +1,122 @@
+"""cadence_static — hold work that changes slower than the hot loop.
+
+A tick pipeline usually runs several cadences at once: a denoise loop
+that repeats every tick, and encoder-side work that only changes when a
+new observation or prompt arrives. Modules on the slower cadence still
+sit inside the hot path, so a graph captures them and pays for them
+every tick even though their inputs are unchanged.
+
+This structure moves such a module out of the loop: its output becomes
+a static buffer the captured graph reads, and the real computation runs
+in an update function the host calls at the module's own cadence. The
+resulting split is explicit — the update callable is returned to the
+caller so it can be registered as a recipe ``outside_update`` rather
+than hidden inside the replacement.
+
+Qualification is empirical: the wrapped output must actually be
+constant across the fast loop. Callers pass calibration captures from
+several iterations of the hot loop and binding refuses when they
+disagree — a module whose output moves per step is not a cadence
+substructure, and freezing it would silently change the model.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Sequence
+
+import torch
+
+from ...guard import PROCEED, GuardedSeam
+
+
+class StaticOutput(GuardedSeam, torch.nn.Module):
+    """Return a buffer the host refreshes at the slower cadence.
+
+    This structure has no input contract to check — ignoring its inputs
+    is the point. Its failure mode is the other one: the host stops
+    calling the update function and the buffer goes stale while every
+    read still succeeds. So the ledger counts reads against refreshes,
+    and a seam read many times and refreshed never is visible as exactly
+    that instead of as a seam that worked.
+    """
+
+    _frt_host_attr = "host_module"
+    _frt_can_fallback = True
+    # the retained host is not a fallback copy: the cadence refresh
+    # projects through it every observation, so its weights must stay
+    # resident — consumption would corrupt the banks silently
+    _frt_host_serving = True
+
+    def __init__(self, original: torch.nn.Module, value: torch.Tensor):
+        super().__init__()
+        self.host_module = original
+        self.register_buffer("buffer", value.contiguous().clone())
+        guard = self._frt_arm(dtypes=None, device=self.buffer.device)
+        guard.notes.update(reads=0, refreshes=0)
+
+    def refreshed(self) -> None:
+        """Record that the slower cadence ran. Called by the updater."""
+        guard = self._frt_guard
+        if guard is not None and not torch.compiler.is_compiling():
+            guard.notes["refreshes"] += 1
+
+    def forward(self, *args, **kwargs):
+        # same rule as the contract check: no Python bookkeeping while a
+        # compiler is tracing this forward
+        if not torch.compiler.is_compiling():
+            self._frt_touch()
+            self._frt_guard.notes["reads"] += 1
+        return self.buffer
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(super().__getattr__("host_module"), name)
+
+
+def bind_cadence_static(
+    modules: Sequence[torch.nn.Module],
+    captures: Sequence[Sequence[torch.Tensor]],
+    *,
+    recompute: Callable[[], Sequence[torch.Tensor]] | None = None,
+    rtol: float = 1e-3,
+    atol: float = 1e-3,
+):
+    """Freeze module outputs into buffers plus one update function.
+
+    ``captures[i]`` holds the outputs module ``i`` produced across
+    several iterations of the hot loop; they must agree, or the module
+    is not on a slower cadence and binding refuses. ``recompute``
+    returns fresh outputs for the same modules — typically by running
+    the host's own encoder path — and the returned update callable
+    copies them into the buffers. Register that callable as the
+    recipe's outside update so the split cadence stays explicit and
+    the slower work is still timed.
+    """
+    if len(modules) != len(captures):
+        raise ValueError("cadence_static: modules/captures mismatch")
+    statics = []
+    for i, (mod, caps) in enumerate(zip(modules, captures)):
+        if not caps:
+            raise ValueError(f"cadence_static: module {i} has no captures")
+        first = caps[0]
+        for other in caps[1:]:
+            if not torch.allclose(first, other, rtol=rtol, atol=atol):
+                raise ValueError(
+                    f"cadence_static: module {i} output varies within "
+                    "the hot loop — not a cadence substructure")
+        statics.append(StaticOutput(mod, first))
+
+    def update() -> None:
+        if recompute is None:
+            raise RuntimeError(
+                "cadence_static: no recompute function was supplied, so "
+                "the buffers cannot be refreshed for a new observation")
+        with torch.no_grad():
+            fresh = recompute()
+            for static, value in zip(statics, fresh):
+                static.buffer.copy_(value)
+                static.refreshed()
+
+    return statics, update

flashrt_structures/impls/cadence_static/cross_attention.py

@@ -0,0 +1,187 @@
+"""Cross-attention K/V addressing for the existing cadence-static structure.
+
+Cross attention has a stable structural signature: Q consumes the current
+hidden width, while K/V consume a different encoder width.  The encoder-side
+K/V projections may therefore be refreshed once at the encoder cadence and
+read from static buffers inside a repeated denoise loop.  Self attention has
+equal Q/K/V input widths and is deliberately excluded.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Mapping, Sequence
+
+import torch
+
+from .buffers import StaticOutput, bind_cadence_static
+
+
+@dataclass(frozen=True)
+class CrossKvCandidate:
+    path: str
+    module: torch.nn.Module
+
+
+def discover_cross_attention_kv(
+    model: torch.nn.Module,
+) -> tuple[CrossKvCandidate, ...]:
+    """Find encoder-side K/V projections without host or model names."""
+    found = []
+    for path, attention in model.named_modules():
+        q_proj = getattr(attention, "to_q", None)
+        k_proj = getattr(attention, "to_k", None)
+        v_proj = getattr(attention, "to_v", None)
+        if not all(
+            isinstance(module, torch.nn.Linear)
+            for module in (q_proj, k_proj, v_proj)
+        ):
+            continue
+        if q_proj.in_features == k_proj.in_features:
+            continue
+        if (
+            k_proj.in_features != v_proj.in_features
+            or k_proj.out_features != v_proj.out_features
+        ):
+            continue
+        found.extend(
+            (
+                CrossKvCandidate(f"{path}.to_k", k_proj),
+                CrossKvCandidate(f"{path}.to_v", v_proj),
+            )
+        )
+    return tuple(found)
+
+
+def capture_cross_attention_kv(
+    candidates: Sequence[CrossKvCandidate],
+    forward: Callable[[], object],
+) -> tuple[tuple[torch.Tensor, ...], ...]:
+    """Capture each candidate across one complete repeated-loop forward."""
+    rows = [[] for _ in candidates]
+    hooks = []
+    for candidate, outputs in zip(candidates, rows):
+        hooks.append(
+            candidate.module.register_forward_hook(
+                lambda _module, _args, output, outputs=outputs:
+                outputs.append(output.detach().clone())
+            )
+        )
+    try:
+        with torch.no_grad():
+            forward()
+    finally:
+        for hook in hooks:
+            hook.remove()
+    return tuple(tuple(outputs) for outputs in rows)
+
+
+def bind_cross_attention_kv(
+    candidates: Sequence[CrossKvCandidate],
+    captures: Sequence[Sequence[torch.Tensor]],
+    *,
+    replacements: Mapping[str, torch.nn.Module] | None = None,
+) -> tuple[dict[str, StaticOutput], tuple[StaticOutput, ...]]:
+    """Bind K/V buffers, optionally consuming already-bound projections."""
+    replacements = replacements or {}
+    modules = []
+    for candidate in candidates:
+        replacement = replacements.get(candidate.path, candidate.module)
+        if getattr(replacement, "_frt_requires_sibling_order", False):
+            # A StashReader (and any equivalent composed tail) is not a
+            # projection in isolation: it reads data produced by a sibling.
+            # Refresh happens outside that sibling call order, so recompute
+            # from the candidate's real projection instead of copying stale
+            # stash contents into the cadence buffer.
+            replacement = candidate.module
+        modules.append(replacement)
+    statics, _ = bind_cadence_static(modules, captures)
+    return (
+        {
+            candidate.path: static
+            for candidate, static in zip(candidates, statics)
+        },
+        tuple(statics),
+    )
+
+
+def refresh_cross_attention_kv(
+    statics: Sequence[StaticOutput],
+    encoder_hidden_states: torch.Tensor,
+) -> None:
+    """Refresh all K/V buffers once before the repeated attention loop."""
+    with torch.no_grad():
+        for static in statics:
+            static.buffer.copy_(
+                static.host_module(encoder_hidden_states)
+            )
+            static.refreshed()
+
+
+def wire_refresh_to_producer(
+    model: torch.nn.Module,
+    statics: Sequence[StaticOutput],
+    forward: Callable[[], object],
+):
+    """Wire the K/V refresh into the producing module's own forward.
+
+    The manual :func:`refresh_cross_attention_kv` form leaves the
+    refresh outside the hot path. That is the right split when only the
+    fast loop is captured — but a *whole-pipeline* capture then records
+    an encoder whose output feeds nothing: the banks are written outside
+    the graph, so replaying on a new observation silently reuses the old
+    encoding. This wires the split shut: one probe forward identifies
+    the module whose output tensor the statics' host projections consume
+    (by object identity), and a forward hook on that producer refreshes
+    every bank whenever it runs. Eager, compiled and captured forms all
+    carry the observation through; within one call the banks are still
+    written once and read every loop step, so the cadence saving stands.
+
+    Returns ``(producer, handle)``; ``handle.remove()`` unwires.
+    Raises ``ValueError`` when no single producer can be identified —
+    the caller keeps the explicit-refresh contract in that case.
+    """
+    if not statics:
+        raise ValueError("cadence_static: no statics to wire")
+    consumed: dict[int, None] = {}
+    probes = []
+    for static in statics:
+        def grab(_module, args, _consumed=consumed):
+            if args and torch.is_tensor(args[0]):
+                _consumed[id(args[0])] = None
+        probes.append(static.register_forward_pre_hook(grab))
+    produced: dict[int, torch.nn.Module] = {}
+
+    def note(module, _args, output):
+        if torch.is_tensor(output):
+            # parents fire after children, so an identity-preserving
+            # wrapper chain resolves to its outermost module
+            produced[id(output)] = module
+
+    watchers = [module.register_forward_hook(note)
+                for _, module in model.named_modules()]
+    try:
+        with torch.no_grad():
+            forward()
+    finally:
+        for hook in probes + watchers:
+            hook.remove()
+    producers = {id(produced[x]): produced[x]
+                 for x in consumed if x in produced}
+    if len(producers) != 1:
+        raise ValueError(
+            "cadence_static: could not identify one producer module for "
+            f"the cross-attention statics ({len(producers)} candidate(s) "
+            "matched by tensor identity)")
+    (producer,) = producers.values()
+
+    def refresh(_module, _args, output):
+        if not torch.is_tensor(output):
+            return None
+        with torch.no_grad():
+            for static in statics:
+                static.buffer.copy_(static.host_module(output))
+                static.refreshed()
+        return None
+
+    return producer, producer.register_forward_hook(refresh)

flashrt_structures/impls/chain_elements.py

@@ -0,0 +1,89 @@
+"""Shared chain elements: the parts every region candidate assembles.
+
+A region family owns two things — a structural identifier and an
+assembly recipe. Everything a recipe *uses* that another recipe could
+use too lives here: weight packing, layout equivalences, activation
+checks, cache duck-typing, and the attention ladder. Keeping the
+elements out of any one family is what keeps a family thin enough to
+read as its recipe, and what keeps two hosts' chains assembling the
+same certified parts instead of drifting copies.
+
+- :func:`fp8_weight` — per-tensor static FP8 packing with the scale
+  returned for alpha folding.
+- :func:`interleave_rows` — the rotate-half ↔ adjacent-pair rotation
+  equivalence, applied to projection rows at pack time; attention dot
+  products are invariant under a shared head-dim permutation.
+- :func:`gelu_tanh_like` — a numeric activation check: the host's
+  callable against tanh-GELU, never a class name.
+- :func:`cache_kv` — per-layer K/V access across cache generations.
+- :data:`ATTN_RUNGS` / :func:`attention_rungs` — the attention
+  element ladder: the house CuTe FA4 runtime first (D256 2CTA for
+  single-KV stacks), the FA2 used-keys entry after it. A rung that
+  loads but cannot execute is eliminated by the binder's functional
+  probe at the bound shapes, never by a device list.
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+
+from . import KernelUnavailable, hub_kernel
+
+FP8_MAX = 448.0
+
+ATTN_RUNGS = (("fa4_cute", "flashrt/fa4-cute-runtime", ">=1",
+               "forward_static"),
+              ("fa2_seqused", "flashrt/fa2-seqused-runtime", ">=1",
+               "forward_seqused_static"))
+
+
+def attention_rungs() -> list[tuple[str, object]]:
+    rungs = []
+    for mode, repo, version, symbol in ATTN_RUNGS:
+        try:
+            kern = hub_kernel(repo, version)
+        except KernelUnavailable:
+            continue
+        if hasattr(kern, symbol):
+            rungs.append((mode, kern))
+    return rungs
+
+
+def fp8_weight(w: torch.Tensor) -> tuple[torch.Tensor, float]:
+    w = w.detach().to("cuda", torch.float32)
+    scale = float(w.abs().amax()) / FP8_MAX
+    if scale <= 0.0:
+        scale = 1.0
+    packed = (w / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+    return packed.contiguous(), scale
+
+
+def interleave_rows(w: torch.Tensor, heads: int,
+                    head_dim: int) -> torch.Tensor:
+    """Permute projection rows so adjacent-pair rotation carries the
+    host's rotate-half convention."""
+    half = head_dim // 2
+    w = w.reshape(heads, head_dim, w.shape[-1])
+    out = torch.empty_like(w)
+    out[:, 0::2] = w[:, :half]
+    out[:, 1::2] = w[:, half:]
+    return out.reshape(heads * head_dim, w.shape[-1])
+
+
+def cache_kv(cache, idx: int):
+    layers = getattr(cache, "layers", None)
+    if layers is not None:
+        return layers[idx].keys, layers[idx].values
+    return cache.key_cache[idx], cache.value_cache[idx]
+
+
+def gelu_tanh_like(act: Callable) -> bool:
+    t = torch.linspace(-4, 4, 65, device="cuda", dtype=torch.bfloat16)
+    try:
+        got = act(t)
+    except Exception:       # noqa: BLE001 — a weird act refuses, not kills
+        return False
+    ref = torch.nn.functional.gelu(t.float(), approximate="tanh")
+    return bool(torch.allclose(got.float(), ref, atol=2e-2))