flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/adapters/diffusers_rotary_attention.py

@@ -0,0 +1,239 @@
+"""Rotary Diffusers attention routed through the dense FA2 structure.
+
+The adapter is selected by the processor boundary and module slots it can
+reproduce, not by a model or processor class name.  It covers the common
+video-transformer form where Q/K are normalised before an optional rotary
+embedding and the processor returns a sequence-major attention result.
+"""
+
+from __future__ import annotations
+
+import inspect
+
+import torch
+
+from ..impls.attention_core import bind_dense_attention_best
+
+
+def _compatible_site(module, processor) -> tuple[bool, str]:
+    if not callable(processor):
+        return False, "processor is not callable"
+    try:
+        parameters = inspect.signature(processor.__call__).parameters
+    except (TypeError, ValueError, AttributeError):
+        return False, "processor call signature is not inspectable"
+    if "rotary_emb" not in parameters:
+        return False, "processor has no rotary_emb boundary"
+    required_modules = ("to_q", "to_k", "to_v", "norm_q", "norm_k")
+    for attr in required_modules:
+        if not isinstance(getattr(module, attr, None), torch.nn.Module):
+            return False, f"attention lacks callable slot {attr!r}"
+    try:
+        out_proj, out_drop = module.to_out[0], module.to_out[1]
+    except (AttributeError, IndexError, KeyError, TypeError):
+        return False, "attention lacks the to_out[projection, dropout] slots"
+    if not all(isinstance(part, torch.nn.Module)
+               for part in (out_proj, out_drop)):
+        return False, "attention output slots are not modules"
+    heads = getattr(module, "heads", None)
+    if not isinstance(heads, int) or heads <= 0:
+        return False, "attention lacks a positive integer head count"
+    if getattr(module, "add_k_proj", None) is not None:
+        return False, "added image KV is not yet an executable form"
+    return True, ""
+
+
+def _projections(attn, hidden_states, encoder_hidden_states):
+    context = hidden_states if encoder_hidden_states is None \
+        else encoder_hidden_states
+    if getattr(attn, "fused_projections", False):
+        if getattr(attn, "is_cross_attention", False):
+            query = attn.to_q(hidden_states)
+            key, value = attn.to_kv(context).chunk(2, dim=-1)
+        else:
+            query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+    else:
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(context)
+        value = attn.to_v(context)
+    return query, key, value
+
+
+def _apply_rotary(hidden_states, rotary_emb):
+    if rotary_emb is None:
+        return hidden_states
+    freqs_cos, freqs_sin = rotary_emb
+    x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+    cos = freqs_cos[..., 0::2]
+    sin = freqs_sin[..., 1::2]
+    out = torch.empty_like(hidden_states)
+    out[..., 0::2] = x1 * cos - x2 * sin
+    out[..., 1::2] = x1 * sin + x2 * cos
+    return out.type_as(hidden_states)
+
+
+def _qkv(attn, hidden_states, encoder_hidden_states, rotary_emb):
+    query, key, value = _projections(
+        attn, hidden_states, encoder_hidden_states)
+    query = attn.norm_q(query)
+    key = attn.norm_k(key)
+    query = query.unflatten(2, (attn.heads, -1))
+    key = key.unflatten(2, (attn.heads, -1))
+    value = value.unflatten(2, (attn.heads, -1))
+    query = _apply_rotary(query, rotary_emb)
+    key = _apply_rotary(key, rotary_emb)
+    return query, key, value
+
+
+class _Recorder:
+    def __init__(self, original, rows):
+        self.original = original
+        self.rows = rows
+
+    def __call__(
+        self, attn, hidden_states, encoder_hidden_states=None,
+        attention_mask=None, rotary_emb=None, *args, **kwargs,
+    ):
+        if attention_mask is None:
+            query, key, value = _qkv(
+                attn, hidden_states, encoder_hidden_states, rotary_emb)
+            row = {
+                "q": query.transpose(1, 2).detach(),
+                "key": key.transpose(1, 2).detach(),
+                "value": value.transpose(1, 2).detach(),
+                "mask": None,
+            }
+            if self.rows:
+                first = self.rows[0]
+                expected = tuple(
+                    (tuple(first[name].shape), first[name].dtype)
+                    for name in ("q", "key", "value"))
+                got = tuple(
+                    (tuple(row[name].shape), row[name].dtype)
+                    for name in ("q", "key", "value"))
+                if got != expected:
+                    raise ValueError(
+                        "attention_core rotary: shape or dtype moved within "
+                        f"one calibration call: {expected} -> {got}")
+            else:
+                # Binding needs one real device sample.  Subsequent calls only
+                # qualify the stable signature; retaining every denoise-step
+                # activation would turn calibration length into VRAM usage.
+                self.rows.append(row)
+        else:
+            self.rows.append({"mask": attention_mask.detach()})
+        return self.original(
+            attn, hidden_states, encoder_hidden_states, attention_mask,
+            rotary_emb, *args, **kwargs)
+
+
+class _FlashRTRotaryAttnProcessor:
+    def __init__(self, core, original):
+        self.core = core
+        self.original = original
+
+    def __call__(
+        self, attn, hidden_states, encoder_hidden_states=None,
+        attention_mask=None, rotary_emb=None, *args, **kwargs,
+    ):
+        if attention_mask is not None:
+            return self.original(
+                attn, hidden_states, encoder_hidden_states, attention_mask,
+                rotary_emb, *args, **kwargs)
+        query, key, value = _qkv(
+            attn, hidden_states, encoder_hidden_states, rotary_emb)
+        projection_dtype = query.dtype
+        guard = getattr(self.core, "_frt_guard", None)
+        accepted_dtypes = tuple(getattr(guard, "dtypes", ()) or ())
+        if accepted_dtypes and projection_dtype not in accepted_dtypes:
+            return self.original(
+                attn, hidden_states, encoder_hidden_states, attention_mask,
+                rotary_emb, *args, **kwargs)
+        hidden_states = self.core(
+            query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
+        hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
+        hidden_states = hidden_states.type_as(query)
+        hidden_states = attn.to_out[0](hidden_states)
+        return attn.to_out[1](hidden_states)
+
+
+class DiffusersRotaryAttentionAdapter:
+    """Route capability-compatible rotary processors through Hub FA2."""
+
+    __name__ = "diffusers_rotary_attention"
+
+    def __call__(self, model, forward, *, prefix_cadence: bool = False):
+        del prefix_cadence
+        sites = []
+        for path, module in model.named_modules():
+            processor = getattr(module, "processor", None)
+            compatible, _ = _compatible_site(module, processor)
+            if compatible:
+                sites.append((path, module, processor))
+        if not sites:
+            return None
+
+        captures = [[] for _ in sites]
+        for (_, module, original), rows in zip(sites, captures):
+            module.processor = _Recorder(original, rows)
+        try:
+            with torch.no_grad():
+                forward()
+        finally:
+            for _, module, original in sites:
+                module.processor = original
+
+        refused = []
+        routes = []
+        observed = {}
+        variants = {}
+        for (path, module, original), rows in zip(sites, captures):
+            if not rows:
+                refused.append((
+                    f"{path}.processor",
+                    "attention_core rotary: compatible processor was not "
+                    "called during calibration",
+                ))
+                continue
+            if any(row.get("mask") is not None for row in rows):
+                refused.append((
+                    f"{path}.processor",
+                    "attention_core rotary: live masks are outside the "
+                    "unmasked executable form",
+                ))
+                continue
+            core = bind_dense_attention_best(rows)
+            if core is None:
+                refused.append((
+                    f"{path}.processor",
+                    "attention_core rotary: Hub FA2 does not cover the "
+                    "captured head dimension",
+                ))
+                continue
+            routed = _FlashRTRotaryAttnProcessor(core, original)
+            routes.append((module, original, routed))
+            observed[f"{path}.processor::fa2_core"] = core
+            variants[f"{path}.processor"] = {
+                "bound": getattr(core, "_frt_variant", "fa2"),
+                "superseded": list(
+                    getattr(core, "_frt_variant_trail", ())),
+            }
+        if not routes:
+            return {}, None, {"refused": refused}
+
+        def enable():
+            for module, _, routed in routes:
+                module.processor = routed
+
+        def disable():
+            for module, original, _ in routes:
+                module.processor = original
+
+        enable()
+        return {}, None, {
+            "revert": [disable],
+            "observed": observed,
+            "toggle": (enable, disable),
+            "refused": refused,
+            "attention_variants": variants,
+        }

flashrt_structures/adapters/factored_qk_norm_rope.py

@@ -0,0 +1,252 @@
+"""Q/K norm + RoPE adapter for factored two-way attention hosts.
+
+The host capability is two independent sibling-QKV projection groups over a
+factored causal/full pack.  Each group is consumed by per-head RMSNorm and the
+same pre-expanded rotate-half position table before a factored attention
+processor.  This is the Cosmos/MoT form, but the adapter deliberately matches
+those slots and dataflow rather than a model or class name.
+"""
+
+from __future__ import annotations
+
+import types
+
+from ..guard import GuardRefused
+from ..impls.qk_norm_rope import bind_per_head_gqa_qk_norm_rope
+from ..impls.qkv_pack.fp8_static import PackedLinear, StashReader
+
+
+_PATHS = (
+    (
+        ("to_q", "to_k", "to_v"),
+        ("norm_q", "norm_k"),
+        "causal_seq",
+    ),
+    (
+        ("add_q_proj", "add_k_proj", "add_v_proj"),
+        ("norm_added_q", "norm_added_k"),
+        "full_only_seq",
+    ),
+)
+
+
+def _pack_parts(plan, path: str, attrs: tuple[str, str, str]):
+    head = plan.swaps.get(f"{path}.{attrs[0]}")
+    key = plan.swaps.get(f"{path}.{attrs[1]}")
+    value = plan.swaps.get(f"{path}.{attrs[2]}")
+    if not (
+        isinstance(head, PackedLinear)
+        and isinstance(key, StashReader)
+        and isinstance(value, StashReader)
+        and key._packed[0] is head
+        and value._packed[0] is head
+    ):
+        return None
+    return head
+
+
+def _epsilon(norm) -> float | None:
+    value = getattr(norm, "variance_epsilon", getattr(norm, "eps", None))
+    return None if value is None else float(value)
+
+
+class FactoredQkNormRopeAdapter:
+    """Compose two packed QKV groups with per-head Q/K norm and RoPE."""
+
+    __name__ = "factored_qk_norm_rope"
+
+    def __call__(self, model, plan):
+        routes = []
+        observed = {}
+        refused = []
+
+        for path, module in model.named_modules():
+            packs = [_pack_parts(plan, path, attrs) for attrs, _, _ in _PATHS]
+            if not any(pack is not None for pack in packs):
+                continue
+            site = f"{path}::factored_qk_norm_rope"
+
+            def refuse(reason: str) -> None:
+                refused.append((site, f"qk_norm_rope refused: {reason}"))
+
+            if not all(pack is not None for pack in packs):
+                refuse("both causal and full QKV groups must be packed")
+                continue
+            if module.training:
+                refuse("training/dropout form is outside the inference seam")
+                continue
+            if not all(
+                hasattr(module, attr)
+                for attr in (
+                    "head_dim",
+                    "num_attention_heads",
+                    "num_key_value_heads",
+                    "dispatch_attention_fn",
+                    "to_out",
+                    "to_add_out",
+                )
+            ):
+                refuse("host lacks the complete factored two-way slots")
+                continue
+            if getattr(module, "cp_mesh", None) is not None:
+                refuse("context-parallel packs are outside this single-device seam")
+                continue
+
+            head_dim = int(module.head_dim)
+            q_heads = int(module.num_attention_heads)
+            kv_heads = int(module.num_key_value_heads)
+            if head_dim != 128:
+                refuse("current Hub entry requires head_dim=128")
+                continue
+
+            bounds = []
+            bad = None
+            for pack, (_, norms, key) in zip(packs, _PATHS):
+                q_norm = getattr(module, norms[0], None)
+                k_norm = getattr(module, norms[1], None)
+                q_weight = getattr(q_norm, "weight", None)
+                k_weight = getattr(k_norm, "weight", None)
+                eps = _epsilon(q_norm)
+                if q_weight is None or k_weight is None or eps is None:
+                    bad = f"{key} Q/K norm weights or epsilon are absent"
+                    break
+                expected = (q_heads * head_dim, kv_heads * head_dim, kv_heads * head_dim)
+                if tuple(pack.splits[:3]) != expected:
+                    bad = f"{key} packed widths {tuple(pack.splits[:3])} != {expected}"
+                    break
+                try:
+                    bound = bind_per_head_gqa_qk_norm_rope(
+                        q_weight,
+                        k_weight,
+                        row_capacity=pack.rows,
+                        q_heads=q_heads,
+                        kv_heads=kv_heads,
+                        head_dim=head_dim,
+                        eps=eps,
+                    )
+                except (ValueError, RuntimeError) as exc:
+                    bad = str(exc)
+                    break
+                bounds.append(bound)
+            if bad is not None:
+                refuse(bad)
+                continue
+
+            original = module.forward
+            had_instance_forward = "forward" in module.__dict__
+            causal_pack, full_pack = packs
+            causal_bound, full_bound = bounds
+
+            def routed(
+                self,
+                pack,
+                attention_mask,
+                packed_position_embeddings,
+                dual_kv_cache=None,
+                natten_metadata=None,
+                *,
+                und_packed=causal_pack,
+                gen_packed=full_pack,
+                und_bound=causal_bound,
+                gen_bound=full_bound,
+            ):
+                del attention_mask
+                if dual_kv_cache is not None:
+                    raise GuardRefused(
+                        "qk_norm_rope: factored cache mutation is outside the bound seam"
+                    )
+                if natten_metadata is not None:
+                    raise GuardRefused(
+                        "qk_norm_rope: neighborhood attention is outside the bound seam"
+                    )
+                if not isinstance(pack, dict) or not all(
+                    key in pack for key in ("causal_seq", "full_only_seq")
+                ):
+                    raise GuardRefused(
+                        "qk_norm_rope: expected a causal/full factored pack"
+                    )
+                try:
+                    cos_pack, sin_pack = packed_position_embeddings
+                    und_cos = cos_pack["causal_seq"]
+                    und_sin = sin_pack["causal_seq"]
+                    gen_cos = cos_pack["full_only_seq"]
+                    gen_sin = sin_pack["full_only_seq"]
+                except (KeyError, TypeError, ValueError) as exc:
+                    raise GuardRefused(
+                        "qk_norm_rope: position tables do not share the factored layout"
+                    ) from exc
+
+                q_und, k_und, v_und = und_bound(
+                    und_packed.joint(pack["causal_seq"]).unsqueeze(0),
+                    und_cos.unsqueeze(0),
+                    und_sin.unsqueeze(0),
+                )
+                q_gen, k_gen, v_gen = gen_bound(
+                    gen_packed.joint(pack["full_only_seq"]).unsqueeze(0),
+                    gen_cos.unsqueeze(0),
+                    gen_sin.unsqueeze(0),
+                )
+                q_und, k_und, v_und = q_und[0], k_und[0], v_und[0]
+                q_gen, k_gen, v_gen = q_gen[0], k_gen[0], v_gen[0]
+
+                if bool(getattr(getattr(self, "config", None), "freeze_und", False)):
+                    q_und = q_und.detach()
+                    k_und = k_und.detach()
+                    v_und = v_und.detach()
+
+                query = dict(pack)
+                key = dict(pack)
+                value = dict(pack)
+                query["causal_seq"], query["full_only_seq"] = q_und, q_gen
+                key["causal_seq"], key["full_only_seq"] = k_und, k_gen
+                value["causal_seq"], value["full_only_seq"] = v_und, v_gen
+                attended = self.dispatch_attention_fn(query, key, value)
+
+                out = dict(pack)
+                out["causal_seq"] = self.to_out(attended["causal_seq"])
+                out["full_only_seq"] = self.to_add_out(attended["full_only_seq"])
+                return out
+
+            routes.append(
+                (
+                    module,
+                    packs,
+                    types.MethodType(routed, module),
+                    original,
+                    had_instance_forward,
+                )
+            )
+            observed[f"{path}.causal::per_head_qk_norm_rope"] = causal_bound
+            observed[f"{path}.full::per_head_qk_norm_rope"] = full_bound
+
+        if not routes:
+            return {"refused": refused} if refused else None
+
+        def enable() -> None:
+            for module, packs, routed, _, _ in routes:
+                for pack in packs:
+                    pack.enable_joint(3)
+                module.forward = routed
+
+        def disable() -> None:
+            for module, packs, _, original, _ in routes:
+                module.forward = original
+                for pack in packs:
+                    pack.disable_joint()
+
+        def revert() -> None:
+            for module, packs, _, original, had_instance_forward in routes:
+                for pack in packs:
+                    pack.disable_joint()
+                if had_instance_forward:
+                    module.forward = original
+                elif "forward" in module.__dict__:
+                    del module.forward
+
+        enable()
+        return {
+            "observed": observed,
+            "revert": [revert],
+            "toggle": (enable, disable),
+            "refused": refused,
+        }

flashrt_structures/adapters/factored_two_way_attention.py

@@ -0,0 +1,99 @@
+"""Adapter for hosts exposing a factored two-way attention processor."""
+
+from __future__ import annotations
+
+from ..impls.attention_core import bind_two_way_attention
+
+
+class FactoredTwoWayAttentionAdapter:
+    """Route compatible factored attention processors through Hub FA2."""
+
+    __name__ = "factored_two_way_attention"
+
+    def __call__(self, model, forward, *, prefix_cadence: bool = False):
+        del prefix_cadence
+        sites = []
+        for path, module in model.named_modules():
+            processor = getattr(module, "dispatch_attention_fn", None)
+            if callable(processor):
+                sites.append((path, module, processor))
+        if not sites:
+            return None
+
+        captures: list[list[dict]] = [[] for _ in sites]
+        for index, (_, module, original) in enumerate(sites):
+            def record(query, key, value, *, i=index, fn=original):
+                captures[i].append(
+                    {"query": query, "key": key, "value": value})
+                return fn(query, key, value)
+
+            module.dispatch_attention_fn = record
+        try:
+            with __import__("torch").no_grad():
+                forward()
+        finally:
+            for _, module, original in sites:
+                module.dispatch_attention_fn = original
+
+        if not any(captures):
+            return None
+        if any(not rows for rows in captures):
+            raise ValueError(
+                "attention_core two_way: only some discovered processors "
+                "were called")
+
+        cores = []
+        routes = []
+        for (path, module, original), rows in zip(sites, captures):
+            first = rows[0]
+
+            def shape(pack):
+                return (
+                    tuple(pack["causal_seq"].shape),
+                    tuple(pack["full_only_seq"].shape),
+                )
+
+            expected = (
+                shape(first["query"]),
+                shape(first["key"]),
+                shape(first["value"]),
+            )
+            for row in rows[1:]:
+                got = (
+                    shape(row["query"]),
+                    shape(row["key"]),
+                    shape(row["value"]),
+                )
+                if got != expected:
+                    raise ValueError(
+                        "attention_core two_way: processor shapes move "
+                        f"within one calibration call: {expected} -> {got}")
+            core = bind_two_way_attention(first)
+
+            def routed(query, key, value, *, bound=core):
+                return bound(query, key, value)
+
+            cores.append((path, core))
+            routes.append((module, original, routed))
+
+        def enable() -> None:
+            for module, _, routed in routes:
+                module.dispatch_attention_fn = routed
+
+        def disable() -> None:
+            for module, original, _ in routes:
+                module.dispatch_attention_fn = original
+
+        def revert() -> None:
+            disable()
+
+        enable()
+        observed = {
+            f"{path}.dispatch_attention_fn::fa2_core": core
+            for path, core in cores
+        }
+        return {}, None, {
+            "revert": [revert],
+            "observed": observed,
+            "toggle": (enable, disable),
+        }