Verified structures: catalog, gates, transactional attach, and provider export

flash_rt/structures/__init__.py

@@ -0,0 +1,70 @@
+"""FlashRT structures — verified, composable model sub-blocks.
+
+A structure is a versioned specification of one model region: boundary
+tensors, framework-neutral weight slots, a plain reference implementation
+used as ground truth, and qualification gates. This package hosts the
+structure catalog and its registry. Implementations, host adapters, and
+the qualification harness build on top of these specifications.
+"""
+
+from flash_rt.structures.registry import StructureSpec, list_structures, load
+
+
+def get(name):
+    """Explicit door: pull one structure and bind it yourself.
+
+    Mirrors ``kernels.get_kernel``: ``get("decoder_ffn").bind(module,
+    calibration=[...])`` returns a gated drop-in replacement you plug in
+    where you choose. See :mod:`flash_rt.structures.handle`.
+    """
+    from flash_rt.structures.handle import get as _get
+
+    return _get(name)
+
+
+def capture(fn, **kwargs):
+    """Capture door: graph a hot stage with declared swap windows.
+
+    See :func:`flash_rt.structures.stages.capture`.
+    """
+    from flash_rt.structures.stages import capture as _capture
+
+    return _capture(fn, **kwargs)
+
+
+def auto_swaps(model, forward, **kwargs):
+    """Distribution layer: discover, calibrate, bind — one pass, no
+    per-seam scaffolding. Returns an :class:`AutoPlan` of swaps.
+
+    See :func:`flash_rt.structures.autobuild.auto_swaps`.
+    """
+    from flash_rt.structures.autobuild import auto_swaps as _auto
+
+    return _auto(model, forward, **kwargs)
+
+
+def run_recipe(recipe, model, ctx=None, **kwargs):
+    """Recipe door: assemble declared levers, audit same-process on the
+    graph, certify or refuse — one call, one receipt.
+
+    See :mod:`flash_rt.structures.recipe` for ``Recipe``/``Lever``/
+    ``Gates`` and the switch lifecycle.
+    """
+    from flash_rt.structures.recipe import run_recipe as _run
+
+    return _run(recipe, model, ctx, **kwargs)
+
+
+def attach(model, forward, **kwargs):
+    """One-call front door: discover, calibrate, gate, activate.
+
+    See :func:`flash_rt.structures.frontdoor.attach`. Imported lazily so
+    that spec-only consumers do not pay for torch-side machinery.
+    """
+    from flash_rt.structures.frontdoor import attach as _attach
+
+    return _attach(model, forward, **kwargs)
+
+
+__all__ = ["StructureSpec", "attach", "capture", "get", "list_structures",
+           "load", "run_recipe"]

flash_rt/structures/adapters/__init__.py

@@ -0,0 +1,13 @@
+"""Host-family adapters — where a structure's seam is host-specific.
+
+Importing this package registers the built-in adapters with autobuild.
+Attention seams (attention_core) live here because where the attention
+math runs differs by host family; a static module pattern cannot find
+them, so each family gets a small adapter.
+"""
+from ..autobuild import register_attention_adapter
+from .gemma_attention import GemmaAttentionAdapter
+
+register_attention_adapter(GemmaAttentionAdapter())
+
+__all__ = ["GemmaAttentionAdapter"]

flash_rt/structures/adapters/gemma_attention.py

@@ -0,0 +1,113 @@
+"""Attention adapter for Gemma-family denoise hosts (pi05 / pi_gemma).
+
+Where the attention math runs is host-specific. In this family the
+transformer's own forward calls ``modeling_gemma.eager_attention_forward``
+directly, bypassing the config/interface dispatch entirely, so the seam
+is that function, not a module. This adapter locates it by capturing one
+denoise pass, binds an :mod:`..impls.attention_core` per layer from the
+captured shapes and masks, and installs a function-level patch that
+routes the fixed denoise shape to the packed-KV kernel while leaving
+prefill and any other shape on the host path.
+
+Registering this adapter lets ``autobuild`` pick up the attention_core
+structure for this host family with no per-host scaffolding at the call
+site — the host still just calls ``auto_swaps``.
+"""
+
+from __future__ import annotations
+
+from ..impls.attention_core import bind_attention_core
+
+
+class GemmaAttentionAdapter:
+    """Recognise a Gemma-family denoise host and wire its fa2 seam."""
+
+    __name__ = "gemma_attention"
+
+    def __call__(self, model, forward):
+        try:
+            import transformers.models.gemma.modeling_gemma as mg
+        except ImportError:
+            return None
+        orig = mg.eager_attention_forward
+
+        recs = {"q": None, "masks": [], "keys": [], "values": []}
+
+        def record(module, query, key, value, attention_mask, **kw):
+            if query.shape[2] < 128:      # denoise (short) vs prefill
+                recs["q"] = query.detach()
+                recs["masks"].append(
+                    attention_mask.detach()
+                    if attention_mask is not None else None)
+                recs["keys"].append(key.detach().clone())
+                recs["values"].append(value.detach().clone())
+            return orig(module, query, key, value, attention_mask, **kw)
+
+        mg.eager_attention_forward = record
+        try:
+            with __import__("torch").no_grad():
+                forward()
+        finally:
+            mg.eager_attention_forward = orig
+        if recs["q"] is None:
+            return None      # host never called this seam — not our family
+
+        n_layers = _infer_layers(model)
+        if n_layers == 0 or len(recs["keys"]) % n_layers != 0:
+            return None
+        steps = len(recs["keys"]) // n_layers
+        captures = [{
+            "q": recs["q"],
+            "keys": [recs["keys"][i + s * n_layers] for s in range(steps)],
+            "values": [recs["values"][i + s * n_layers]
+                       for s in range(steps)],
+            "mask": recs["masks"][i],
+        } for i in range(n_layers)]
+
+        bound = bind_attention_core(captures)
+        if bound is None:
+            return None      # head_dim unsupported → host keeps its path
+        cores, prefix_update = bound
+        seq_q = recs["q"].shape[2]
+        expert = _expert_layers(model)
+        for i, layer in enumerate(expert):
+            layer.self_attn._fa2_core = cores[i]
+
+        fired = {"n": 0}
+
+        def fa2_fn(module, query, key, value, attention_mask, **kw):
+            if query.shape[2] != seq_q or not hasattr(module, "_fa2_core"):
+                return orig(module, query, key, value, attention_mask, **kw)
+            fired["n"] += 1
+            return module._fa2_core(query, key, value,
+                                    scale=kw.get("scaling")), None
+
+        mg.eager_attention_forward = fa2_fn
+        # self-verify the patch actually takes over the seam on the same
+        # forward that will be captured; if it never fires (the host's
+        # captured entry uses a different attention shape than the one
+        # we calibrated on) leave the host untouched rather than install
+        # a patch that only adds a Python check and blocks fusion
+        with __import__("torch").no_grad():
+            forward()
+        if fired["n"] == 0:
+            mg.eager_attention_forward = orig
+            for layer in expert:
+                if hasattr(layer.self_attn, "_fa2_core"):
+                    del layer.self_attn._fa2_core
+            return None
+        # the swap map is empty (the seam is a function, not a module);
+        # the patch and the per-layer core buffers are the swap
+        return {}, None
+
+
+def _infer_layers(model) -> int:
+    layers = _expert_layers(model)
+    return len(layers) if layers is not None else 0
+
+
+def _expert_layers(model):
+    try:
+        return model.paligemma_with_expert.gemma_expert.model.layers
+    except AttributeError:
+        return None