flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/adarms_stack/region.py

@@ -0,0 +1,102 @@
+"""The adarms_stack region family: structural identification, candidates.
+
+The identifier matches shape, never names: a module carrying a
+``layers`` list whose blocks each hold a bias-free attention group
+(``q_proj``/``k_proj``/``v_proj``/``o_proj``), a bias-free gated FFN
+(``gate_proj``/``up_proj``/``down_proj``), and *conditioned* norms — a
+``dense`` projection emitting three modulation vectors per norm is
+what separates this stack from an ordinary decoder tower — plus the
+stack-level tail (a conditioned final norm and a rotary table). The
+key width must be narrower than the query width: the fused chain's
+cache layout is written for the single-KV band, and a stack outside
+it is exactly what the candidate's own fact checks exist to refuse.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from . import fp8_chain
+from ... import regions
+
+
+def _plain_linear(mod) -> bool:
+    return isinstance(mod, torch.nn.Linear) and mod.bias is None
+
+
+def _ada_norm(mod, dim: int) -> bool:
+    dense = getattr(mod, "dense", None)
+    return (isinstance(dense, torch.nn.Linear)
+            and dense.out_features == 3 * dim)
+
+
+def _block_ok(block, dim: int) -> bool:
+    attn = getattr(block, "self_attn", None)
+    mlp = getattr(block, "mlp", None)
+    if attn is None or mlp is None:
+        return False
+    if not all(_plain_linear(getattr(attn, a, None))
+               for a in ("q_proj", "k_proj", "v_proj", "o_proj")):
+        return False
+    if not all(_plain_linear(getattr(mlp, a, None))
+               for a in ("gate_proj", "up_proj", "down_proj")):
+        return False
+    if mlp.gate_proj.out_features != mlp.up_proj.out_features:
+        return False
+    if not _ada_norm(getattr(block, "input_layernorm", None), dim):
+        return False
+    return _ada_norm(getattr(block, "post_attention_layernorm", None),
+                     dim)
+
+
+def identify(model) -> list[str]:
+    roots = []
+    for path, mod in model.named_modules():
+        layers = getattr(mod, "layers", None)
+        if not isinstance(layers, torch.nn.ModuleList) or len(layers) < 2:
+            continue
+        if not callable(getattr(mod, "rotary_emb", None)):
+            continue
+        head = layers[0]
+        attn = getattr(head, "self_attn", None)
+        if attn is None or not isinstance(
+                getattr(attn, "q_proj", None), torch.nn.Linear):
+            continue
+        dim = attn.q_proj.in_features
+        if not _ada_norm(getattr(mod, "norm", None), dim):
+            continue
+        if not all(_block_ok(b, dim) for b in layers):
+            continue
+        if attn.k_proj.out_features >= attn.q_proj.out_features:
+            continue
+        roots.append(path)
+    return roots
+
+
+def _band_candidate(band: str, row: dict) -> regions.RegionCandidate:
+    return regions.RegionCandidate(
+        name=f"{band}_chain",
+        missing=lambda band=band: fp8_chain.missing_symbols(band=band),
+        bind=lambda model, root, probe, band=band:
+            fp8_chain.bind_adarms_fp8_chain(model, root, probe,
+                                            band=band),
+        precision_rank=row["precision_rank"],
+    )
+
+
+#: the candidate list is generated from the band table — adding a
+#: precision band to the family is a table row, not new wiring here
+FAMILY = regions.RegionFamily(
+    family="adarms_stack",
+    identify=identify,
+    candidates=[_band_candidate(band, row)
+                for band, row in fp8_chain.BANDS.items()],
+)
+
+
+def register() -> None:
+    """(Re-)register the family — idempotent, import calls it once."""
+    regions.register_region_family(FAMILY)
+
+
+register()

flashrt_structures/impls/attention_core/__init__.py

@@ -0,0 +1,132 @@
+import torch
+
+from .fa2_seqused import (DenseAttention, PackedKVAttention,
+                          bind_attention_core,
+                          bind_dense_attention, plan_packed_kv)
+from .two_way_fa2 import FactoredTwoWayAttention, bind_two_way_attention
+
+
+def bind_dense_attention_best(captures):
+    """Dense attention across the variant family, precision-descending.
+
+    One structure, parallel executable forms, and no second hardware
+    table: each variant's kernel package declares its own archs (or
+    ships build variants only for the hosts it serves), so the device
+    split *is* the refusal machinery. The order is precision first:
+    FA2 (BF16, numerics preserving) binds wherever its runtime
+    executes — every current receipt keeps its exact form. Then the
+    FA4 CuTe DSL form (BF16, the SM100-family hot path), the
+    allocation-free masked MHA (BF16/FP16, batch-of-one sites), and
+    last the FP8 FA4 form. All are judged by the same downstream
+    gates; a device none serves keeps the host's own attention.
+
+    The winner carries the trail: ``_frt_variant`` names the bound
+    form and ``_frt_variant_trail`` holds what each preferred variant
+    said when it stepped aside. Without that, a host whose preferred
+    package is merely absent looks identical to one where it was
+    weighed and rejected — the two need different fixes, and only the
+    trail tells them apart.
+    """
+    # a package can refuse a device three ways: its arch declaration
+    # (ValueError from the loader's metadata check), the kernels
+    # library finding no build variant for the host (OSError), or a
+    # bind smoke the runtime cannot execute (RuntimeError) — each
+    # means "not this variant here", never an error for the family
+    def _fa2(caps):
+        return bind_dense_attention(caps)
+
+    def _fa4_cute(caps):
+        from . import fa4_cute
+        return fa4_cute.bind_dense_attention(caps)
+
+    def _masked_mha(caps):
+        from . import masked_mha
+        return masked_mha.bind_dense_attention(caps)
+
+    def _fa4_fp8(caps):
+        from . import fa4_fp8
+        return fa4_fp8.bind_dense_attention(caps)
+
+    refusals, declined = [], 0
+    for name, binder in (("fa2", _fa2), ("fa4_cute", _fa4_cute),
+                         ("masked_mha", _masked_mha),
+                         ("fa4_fp8", _fa4_fp8)):
+        try:
+            core = binder(captures)
+        except (ValueError, RuntimeError, OSError) as refusal:
+            refusals.append(f"{name}: {str(refusal)[:120]}")
+            continue
+        if core is not None and not _beats_host(core, captures[0]):
+            refusals.append(
+                f"{name}: bound but measured slower than the host "
+                "attention at the captured shape — stepped aside")
+            declined += 1
+            core = None
+        if core is not None:
+            # the seam is served, but which variant served it and what
+            # the preferred ones said are both load-bearing facts: a
+            # host silently falling back to a lower-precision or slower
+            # form is exactly the failure the ordering exists to make
+            # visible, and it is invisible unless the superseded
+            # refusals travel with the bound module
+            core._frt_variant = name
+            core._frt_variant_trail = tuple(refusals)
+            return core
+        declined += 1
+        refusals.append(f"{name}: declined the captured shape form")
+    if declined:
+        # at least one variant executed its qualification and declined
+        # the shape form — a site-level refusal the adapter records,
+        # same contract as a single binder returning None
+        return None
+    raise ValueError(
+        "attention_core: no variant serves this device — "
+        + "; ".join(refusals))
+
+
+def _beats_host(core, capture, margin: float = 0.02, iters: int = 20):
+    """The family's speed gate: a variant seats only if it measures at
+    least as fast as the host's own attention on the captured shape.
+    Availability and precision order decide who gets weighed first;
+    this decides whether the winner actually serves — bands are
+    measured, not conceded, in this family too."""
+    import torch.nn.functional as F
+
+    query = capture.get("q")
+    key = capture.get("key", capture.get("k"))
+    value = capture.get("value", capture.get("v"))
+    mask = capture.get("mask")
+    if query is None or key is None or value is None:
+        return True
+    if not query.is_cuda:
+        return True
+
+    def _time(fn):
+        with torch.no_grad():
+            for _ in range(4):
+                fn()
+            torch.cuda.synchronize()
+            start = torch.cuda.Event(True)
+            end = torch.cuda.Event(True)
+            start.record()
+            for _ in range(iters):
+                fn()
+            end.record()
+            torch.cuda.synchronize()
+        return start.elapsed_time(end) / iters
+
+    try:
+        ours = _time(lambda: core(query, key, value))
+        host = _time(lambda: F.scaled_dot_product_attention(
+            query, key, value, attn_mask=mask))
+    except (RuntimeError, ValueError):
+        return False
+    return ours <= host * (1.0 + margin)
+
+
+__all__ = ["DenseAttention", "PackedKVAttention",
+           "bind_attention_core", "bind_dense_attention",
+           "bind_dense_attention_best", "plan_packed_kv",
+           "FactoredTwoWayAttention", "bind_two_way_attention"]
+# variant modules (fa4_fp8, fa4_cute, masked_mha) import lazily inside
+# the family binder: loading one must not require the others' runtimes

flashrt_structures/impls/attention_core/fa2_seqused.py

@@ -0,0 +1,458 @@
+"""attention_core — run attention on the FlashRT FA2 kernel.
+
+A fused attention kernel wants contiguous keys and values. This module
+provides two executable forms: a cadence-aware packed-KV form for decoder
+loops, and a stateless dense form for hosts that supply complete Q/K/V on
+every call. Both reduce supported masks to contiguous allowed key ranges
+before dispatching one dense attention call.
+
+The packed prefix (everything before the blocked run) belongs to a
+slower cadence — it is the encoder's output for the current
+observation, unchanged across the denoise loop — so it is filled once
+at bind time and refreshed through an update function, exactly as
+:mod:`..cadence_static` does for whole modules. The suffix is written
+per call.
+
+Three qualifications, all decided from real captures:
+
+- ``head_dim`` must be one the kernel supports; otherwise the caller
+  keeps its own path (``bind_attention_core`` returns ``None`` so the
+  host can fall back to a community kernel rather than fail),
+- every query row must see the same mask, and the blocked positions
+  must form a single contiguous run — anything else is not expressible
+  as a packed dense attention,
+- the prefix keys must not move across the loop; a moving prefix means
+  the split is wrong and the parity gate would catch it downstream, so
+  it is rejected here where the reason is still legible.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardedSeam
+
+@lru_cache(maxsize=1)
+def supported_head_dims() -> tuple[int, ...]:
+    """Read the executable envelope from the installed Hub artifact."""
+    package = hub_kernel("flashrt/fa2-seqused-runtime", ">=1")
+    advertised = getattr(package, "SUPPORTED_HEAD_DIMS", None)
+    if advertised is None:
+        raise ValueError(
+            "attention_core: FA2 Hub artifact does not advertise "
+            "SUPPORTED_HEAD_DIMS; refusing to duplicate backend capability "
+            "inside the structure layer")
+    dims = tuple(sorted({int(dim) for dim in advertised}))
+    if not dims or any(dim <= 0 for dim in dims):
+        raise ValueError(
+            "attention_core: FA2 Hub artifact advertised invalid head dims")
+    return dims
+
+
+@dataclass
+class _Scratch:
+    """Output/LSE/workspace shared by same-shaped attention sites."""
+
+    out: torch.Tensor
+    lse: torch.Tensor
+    workspace: object
+
+
+@dataclass
+class PackedKVPlan:
+    """How the host's masked attention maps onto a dense one."""
+
+    prefix: int
+    suffix_start: int
+    suffix_len: int
+    seq_kv: int
+
+    @property
+    def packed(self) -> int:
+        return self.prefix + self.suffix_len
+
+
+def plan_packed_kv(mask: torch.Tensor | None, seq_kv: int) -> PackedKVPlan:
+    """Derive the packing plan from one captured attention mask."""
+    if mask is None:
+        return PackedKVPlan(seq_kv, seq_kv, 0, seq_kv)
+    if mask.dim() < 3:
+        raise ValueError("attention_core: unexpected mask rank")
+    rows = mask.reshape(-1, mask.shape[-2], mask.shape[-1])[0]
+    if not bool((rows == rows[0]).all()):
+        raise ValueError("attention_core: mask differs per query row")
+    row = rows[0].float()
+    blocked = ((row < -1e5) | row.isneginf()).nonzero().flatten()
+    if blocked.numel() == 0:
+        return PackedKVPlan(seq_kv, seq_kv, 0, seq_kv)
+    runs, start, prev = [], None, None
+    for i in blocked.tolist():
+        if start is None:
+            start = prev = i
+        elif i == prev + 1:
+            prev = i
+        else:
+            runs.append((start, prev))
+            start = prev = i
+    runs.append((start, prev))
+    if len(runs) != 1:
+        raise ValueError(
+            f"attention_core: mask blocks {len(runs)} separate runs — "
+            "not expressible as one packed dense attention")
+    lo, hi = runs[0]
+    return PackedKVPlan(lo, hi + 1, seq_kv - (hi + 1), seq_kv)
+
+
+class PackedKVAttention(GuardedSeam, torch.nn.Module):
+    """Attention over packed keys/values, on the FlashRT FA2 kernel.
+
+    Holds one host module's packed buffers. Call it in place of the
+    host's attention body; the prefix half is refreshed by the update
+    function returned from :func:`bind_attention_core`.
+
+    The packed region, the output and the split-KV workspace are all
+    allocated for one query length, and this module is reached through a
+    routed call rather than a module path, so there is no host module to
+    revert to per call: a query outside the calibrated form raises. That
+    is the whole reason to check — handing the kernel a scratch buffer
+    sized for a different sequence is the failure that does not announce
+    itself.
+    """
+
+    def __init__(self, plan: PackedKVPlan, q_shape, kv_heads: int,
+                 dtype: torch.dtype, device, prefix_kv=None,
+                 scratch: "_Scratch | None" = None):
+        super().__init__()
+        self.plan = plan
+        b, heads, seq_q, head_dim = q_shape
+        self.seq_q = seq_q
+        # set by alias_suffix when a producer writes that side in place
+        self._alias_k = False
+        self._alias_v = False
+        self._kfa = hub_kernel("flashrt/fa2-seqused-runtime", ">=1")
+        self.register_buffer("packed_k", torch.zeros(
+            b, plan.packed, kv_heads, head_dim, device=device,
+            dtype=dtype))
+        self.register_buffer("packed_v", torch.zeros_like(self.packed_k))
+        if prefix_kv is not None:
+            k0, v0 = prefix_kv
+            self.packed_k[:, :plan.prefix] = k0
+            self.packed_v[:, :plan.prefix] = v0
+        # Output, LSE and split-KV workspace are scratch: each site
+        # consumes its result before the next one runs, so sites with
+        # the same shapes share one set. Per-site copies cost real
+        # latency (pi05 r16: 18 private scratches ran 1.5% slower than
+        # one shared set) and buy nothing.
+        if scratch is None:
+            q_sample = torch.empty(b, seq_q, heads, head_dim,
+                                   device=device, dtype=dtype)
+            out, lse = self._kfa.allocate_outputs(q_sample)
+            scratch = _Scratch(out, lse, self._kfa.allocate_workspace(
+                q_sample, self.packed_k))
+        self._scratch = scratch
+        self._frt_arm(dtypes=(dtype,), device=self.packed_k.device,
+                      k=int(head_dim), rows=int(b * heads * seq_q))
+
+    def forward(self, query, key, value, *, scale=None):
+        """``query``/``key``/``value`` in the host's (B, H, S, D)."""
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:            # unreachable: no host path
+            return admitted                    # to revert to, so it raises
+        plan = self.plan
+        q = query.transpose(1, 2).contiguous()
+        if plan.suffix_len:
+            self.packed_k[:, plan.prefix:].copy_(
+                key.transpose(1, 2)[:, plan.suffix_start:])
+            self.packed_v[:, plan.prefix:].copy_(
+                value.transpose(1, 2)[:, plan.suffix_start:])
+        sc = self._scratch
+        return self._kfa.forward_static(
+            q, self.packed_k, self.packed_v, out=sc.out,
+            softmax_lse=sc.lse, workspace=sc.workspace,
+            softmax_scale=scale)
+
+    def alias_suffix(self, *, key: bool = False, value: bool = False):
+        """Hand out the packed suffix rows for a producer to write into.
+
+        The declarative form of what a hand-written runtime does when it
+        gives the next stage a pointer: instead of the producer filling
+        its own buffer and this module copying it in, the producer's
+        output *is* the region.
+
+        Only legal where nothing transforms the tensor between the two —
+        a rotary embedding applied after the projection would leave the
+        untransformed values here. Callers therefore alias key and value
+        independently, and take ``None`` for whatever does not qualify.
+        """
+        plan = self.plan
+        if not plan.suffix_len or self.packed_k.shape[0] != 1:
+            return None, None          # only a single batch row slices
+        regions = []                   # into a contiguous suffix
+        for want, packed, flag in ((key, self.packed_k, "_alias_k"),
+                                   (value, self.packed_v, "_alias_v")):
+            if not want:
+                regions.append(None)
+                continue
+            region = packed[0, plan.prefix:]
+            if not region.is_contiguous():
+                regions.append(None)
+                continue
+            setattr(self, flag, True)
+            regions.append(region)
+        return tuple(regions)
+
+    def forward_suffix(self, query, key, value, *, scale=None):
+        """Kernel layout (B, S, H, D), with **suffix-only** keys/values.
+
+        Two things separate this from :meth:`forward`. The layout: the
+        host builds (B, H, S, D) because that is what eager SDPA wants
+        and :meth:`forward` transposes it back — two cancelling
+        transposes plus the copies that make each of them contiguous. A
+        caller that owns the whole attention sublayer never builds the
+        host layout at all; the projections' own output view *is* this
+        one. And the extent: the host hands attention the full
+        cache-concatenated KV and the suffix gets sliced back out here,
+        while a sublayer's projections produce exactly the new tokens —
+        which is the suffix already. Binding checks that equality rather
+        than assuming it.
+        """
+        # the contract is the same one ``forward`` checks: it counts rows
+        # against the head dim, which is what both layouts agree on
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        plan = self.plan
+        q = query if query.is_contiguous() else query.contiguous()
+        if plan.suffix_len:
+            # an aliased side already wrote itself here
+            if not self._alias_k:
+                self.packed_k[:, plan.prefix:].copy_(key)
+            if not self._alias_v:
+                self.packed_v[:, plan.prefix:].copy_(value)
+        sc = self._scratch
+        return self._kfa.forward_static(
+            q, self.packed_k, self.packed_v, out=sc.out,
+            softmax_lse=sc.lse, workspace=sc.workspace,
+            softmax_scale=scale)
+
+
+class DenseAttention(GuardedSeam, torch.nn.Module):
+    """FA2 replacement for an ordinary dense SDPA call.
+
+    Unlike :class:`PackedKVAttention`, this form owns no observation-cadence
+    state.  Every call receives the complete Q/K/V tensors, so it is suitable
+    for Diffusers self- and cross-attention and remains correct when the
+    conditioning changes between graph replays.
+
+    Inputs and outputs use the host SDPA layout ``[B, H, S, D]``.  The package
+    consumes ``[B, S, H, D]``; reversing the host's projection view is normally
+    already contiguous and therefore does not materialise a transpose.
+    """
+
+    def __init__(
+        self, q_shape, kv_shape, dtype: torch.dtype, device,
+        allowed_ranges=None, scratch: "_Scratch | None" = None,
+    ):
+        super().__init__()
+        b, heads, seq_q, head_dim = q_shape
+        kb, kv_heads, seq_kv, kv_dim = kv_shape
+        if kb != b or kv_dim != head_dim:
+            raise ValueError(
+                "attention_core dense: Q and KV batch/head dimensions differ")
+        if heads % kv_heads:
+            raise ValueError(
+                "attention_core dense: query heads must be divisible by "
+                "KV heads")
+        self.q_shape = tuple(q_shape)
+        self.kv_shape = tuple(kv_shape)
+        self.allowed_ranges = tuple(allowed_ranges or ())
+        self._kfa = hub_kernel("flashrt/fa2-seqused-runtime", ">=1")
+        q_sample = torch.empty(
+            b, seq_q, heads, head_dim, device=device, dtype=dtype)
+        packed_seq = (
+            sum(hi - lo for lo, hi in self.allowed_ranges)
+            if self.allowed_ranges else seq_kv)
+        kv_sample = torch.empty(
+            b, packed_seq, kv_heads, head_dim, device=device, dtype=dtype)
+        if self.allowed_ranges:
+            self.register_buffer("packed_k", torch.empty_like(kv_sample))
+            self.register_buffer("packed_v", torch.empty_like(kv_sample))
+        if scratch is None:
+            out, lse = self._kfa.allocate_outputs(q_sample)
+            scratch = _Scratch(
+                out, lse, self._kfa.allocate_workspace(q_sample, kv_sample))
+        elif (
+            scratch.out.shape != q_sample.shape
+            or scratch.out.dtype != dtype
+            or scratch.out.device != q_sample.device
+            or scratch.lse.shape != (b, heads, seq_q)
+            or scratch.lse.device != q_sample.device
+        ):
+            raise ValueError(
+                "attention_core dense: shared scratch does not match "
+                "the bound attention form"
+            )
+        self._scratch = scratch
+        self._frt_arm(
+            dtypes=(dtype,), device=q_sample.device, k=int(head_dim),
+            rows=int(b * heads * seq_q))
+
+    def forward(self, query, key, value, *, scale=None):
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        if tuple(query.shape) != self.q_shape:
+            raise ValueError(
+                "attention_core dense: query shape moved from "
+                f"{self.q_shape} to {tuple(query.shape)}")
+        if tuple(key.shape) != self.kv_shape or value.shape != key.shape:
+            raise ValueError(
+                "attention_core dense: key/value shape moved from "
+                f"{self.kv_shape} to {tuple(key.shape)}/{tuple(value.shape)}")
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        if not q.is_contiguous():
+            q = q.contiguous()
+        if self.allowed_ranges:
+            offset = 0
+            for lo, hi in self.allowed_ranges:
+                length = hi - lo
+                self.packed_k[:, offset:offset + length].copy_(k[:, lo:hi])
+                self.packed_v[:, offset:offset + length].copy_(v[:, lo:hi])
+                offset += length
+            k, v = self.packed_k, self.packed_v
+        else:
+            if not k.is_contiguous():
+                k = k.contiguous()
+            if not v.is_contiguous():
+                v = v.contiguous()
+        sc = self._scratch
+        out = self._kfa.forward_static(
+            q, k, v, out=sc.out, softmax_lse=sc.lse,
+            workspace=sc.workspace, softmax_scale=scale)
+        return out.transpose(1, 2)
+
+
+def _allowed_ranges(mask):
+    if mask is None:
+        return ()
+    rows = mask.reshape(-1, mask.shape[-1])
+    first = rows[0]
+    if mask.dtype == torch.bool:
+        allowed = first
+        if not bool((rows == first).all()):
+            raise ValueError(
+                "attention_core dense: mask differs per query row")
+    else:
+        allowed = ~((first.float() < -1e5) | first.float().isneginf())
+        other = ~((rows.float() < -1e5) | rows.float().isneginf())
+        if not bool((other == allowed).all()):
+            raise ValueError(
+                "attention_core dense: mask differs per query row")
+    indices = allowed.nonzero().flatten().tolist()
+    if not indices:
+        raise ValueError("attention_core dense: mask permits no keys")
+    ranges = []
+    for index in indices:
+        if not ranges or index != ranges[-1][1]:
+            ranges.append([index, index + 1])
+        else:
+            ranges[-1][1] += 1
+    if len(ranges) > 8:
+        # the packed-copy loop is linear in segments; past a handful
+        # the copies outweigh the masked-out keys and the masked
+        # executable form serves better
+        return None
+    if len(ranges) == 1 and ranges[0] == [0, mask.shape[-1]]:
+        return ()
+    return tuple((lo, hi) for lo, hi in ranges)
+
+
+def bind_dense_attention(captures):
+    """Bind one stateless dense FA2 core from repeated host captures."""
+    if not captures:
+        raise ValueError("attention_core dense: no captures")
+    first = captures[0]
+    query, key, value = first["q"], first["key"], first["value"]
+    head_dim = query.shape[-1]
+    if head_dim not in supported_head_dims():
+        return None
+    allowed_ranges = _allowed_ranges(first.get("mask"))
+    if allowed_ranges is None:
+        return None
+    expected = (tuple(query.shape), tuple(key.shape), tuple(value.shape),
+                query.dtype, key.dtype, value.dtype)
+    for capture in captures[1:]:
+        got = (
+            tuple(capture["q"].shape),
+            tuple(capture["key"].shape),
+            tuple(capture["value"].shape),
+            capture["q"].dtype,
+            capture["key"].dtype,
+            capture["value"].dtype,
+        )
+        if got != expected:
+            raise ValueError(
+                "attention_core dense: shape, dtype, or mask moved within "
+                f"one calibration call: {expected} -> {got}")
+        if _allowed_ranges(capture.get("mask")) != allowed_ranges:
+            raise ValueError(
+                "attention_core dense: mask pattern moved within one "
+                "calibration call")
+    if not (query.dtype == key.dtype == value.dtype):
+        raise ValueError("attention_core dense: Q/K/V dtypes differ")
+    return DenseAttention(
+        query.shape, key.shape, query.dtype, query.device,
+        allowed_ranges=allowed_ranges)
+
+
+def bind_attention_core(captures, *, prefix_static_rtol: float = 1e-3):
+    """Bind one packed-KV attention per site from real captures.
+
+    ``captures`` is a sequence of per-site dicts holding ``q`` (one
+    captured query, host layout), ``keys``/``values`` (the tensors that
+    site produced across the hot loop, oldest first) and ``mask``.
+    Returns ``(modules, update)``, or ``None`` when the head dim is
+    unsupported so the caller can keep a fallback path.
+    """
+    if not captures:
+        raise ValueError("attention_core: no captures")
+    head_dim = captures[0]["q"].shape[-1]
+    if head_dim not in supported_head_dims():
+        return None
+
+    modules, scratch = [], None
+    for site, cap in enumerate(captures):
+        keys = cap["keys"]
+        plan = plan_packed_kv(cap.get("mask"), keys[0].shape[2])
+        first = keys[0][:, :, :plan.prefix]
+        for other in keys[1:]:
+            if not torch.allclose(first, other[:, :, :plan.prefix],
+                                  rtol=prefix_static_rtol,
+                                  atol=prefix_static_rtol):
+                raise ValueError(
+                    f"attention_core: site {site} prefix keys move "
+                    "across the loop — the cadence split is wrong")
+        q = cap["q"]
+        core = PackedKVAttention(
+            plan, q.shape, keys[0].shape[1], q.dtype, q.device,
+            prefix_kv=(keys[0].transpose(1, 2)[:, :plan.prefix],
+                       cap["values"][0].transpose(1, 2)[:, :plan.prefix]),
+            scratch=scratch)
+        scratch = scratch or core._scratch
+        modules.append(core)
+
+    def update(fresh_kv) -> None:
+        """Refresh every site's prefix from freshly computed K/V."""
+        with torch.no_grad():
+            for mod, (k, v) in zip(modules, fresh_kv):
+                p = mod.plan.prefix
+                mod.packed_k[:, :p].copy_(k.transpose(1, 2)[:, :p])
+                mod.packed_v[:, :p].copy_(v.transpose(1, 2)[:, :p])
+
+    return modules, update