flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/gated_delta_core/hub_v3.py

@@ -0,0 +1,152 @@
+"""Hub v3 executable forms for ``gated_delta_core``."""
+
+from __future__ import annotations
+
+import torch
+
+from ...guard import PROCEED, GuardedSeam
+from .. import hub_kernel
+
+
+class HubV3GatedDeltaCore(GuardedSeam, torch.nn.Module):
+    """Single-token H=32/48, D=128 recurrence with explicit state output.
+
+    The log-decay ``g`` binds in the dtype the host actually exposes:
+    BF16 through the original entry, FP32 through the ``gf32`` twin —
+    the 27B-class cached-decode hosts keep ``g`` in FP32, and rounding
+    it through BF16 (or casting in the hot path) is what qualification
+    used to refuse here.
+    """
+
+    def __init__(self, sample: torch.Tensor,
+                 g_dtype: torch.dtype = torch.bfloat16,
+                 state_dtype: torch.dtype = torch.bfloat16):
+        super().__init__()
+        if sample.dtype != torch.bfloat16:
+            raise ValueError("gated_delta_core v3 requires BF16 Q/K/V")
+        if g_dtype not in (torch.bfloat16, torch.float32):
+            raise ValueError(
+                "gated_delta_core v3 serves BF16 or FP32 log-decay only")
+        if state_dtype not in (torch.bfloat16, torch.float32):
+            raise ValueError(
+                "gated_delta_core v3 serves BF16 or FP32 state only")
+        if state_dtype is torch.float32 and g_dtype is not torch.float32:
+            raise ValueError(
+                "gated_delta_core v3 has no BF16-g/FP32-state entry; no "
+                "host has exposed that combination")
+        self._g_dtype = g_dtype
+        self._state_dtype = state_dtype
+        if sample.ndim != 4 or sample.shape[0] != 1 \
+                or sample.shape[1] != 1 \
+                or sample.shape[2] not in (32, 48) \
+                or sample.shape[3] != 128:
+            raise ValueError(
+                "gated_delta_core v3 requires Q shape "
+                "(1,1,H,128) with H=32 or H=48; the published "
+                "sequence API has no explicit state output")
+        if not sample.is_contiguous():
+            raise ValueError("gated_delta_core v3 requires contiguous Q/K/V")
+        self.heads = int(sample.shape[2])
+        self._ops = hub_kernel("flashrt/gated-delta-attention", ">=3")
+        if g_dtype is torch.float32:
+            name = ("gated_delta_recurrent_inout_gf32_sf32_bf16"
+                    if state_dtype is torch.float32
+                    else "gated_delta_recurrent_inout_gf32_bf16")
+            step = getattr(self._ops, name, None)
+            if step is None:
+                raise ValueError(
+                    "refused: the installed gated-delta-attention build "
+                    f"predates the {name} entry; a release carrying it "
+                    "is required")
+        else:
+            step = self._ops.gated_delta_recurrent_inout_bf16
+        self._step = step
+        self.register_buffer(
+            "_state_out",
+            torch.empty(
+                1, self.heads, 128, 128,
+                device=sample.device, dtype=self._state_dtype),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_out",
+            torch.empty(
+                1, self.heads, 128,
+                device=sample.device, dtype=torch.bfloat16),
+            persistent=False,
+        )
+        self._frt_arm(
+            dtypes=(torch.bfloat16,), device=sample.device, k=128,
+            rows=self.heads)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        log_decay: torch.Tensor,
+        beta: torch.Tensor,
+        state: torch.Tensor | None,
+        *,
+        output_final_state: bool,
+        use_qk_l2norm: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        if query.ndim != 4 or query.shape[0] != 1 \
+                or query.shape[1:] != (1, self.heads, 128):
+            raise ValueError(
+                "gated_delta_core v3 query shape moved after binding")
+        if key.shape != query.shape or value.shape != query.shape:
+            raise ValueError("gated_delta_core v3 Q/K/V shapes differ")
+        if not (query.is_contiguous() and key.is_contiguous()
+                and value.is_contiguous()):
+            raise ValueError("gated_delta_core v3 requires contiguous Q/K/V")
+        if log_decay.shape != query.shape[:3] \
+                or beta.shape != log_decay.shape:
+            raise ValueError("gated_delta_core v3 gating shapes differ")
+        if log_decay.dtype != self._g_dtype \
+                or beta.dtype != torch.bfloat16:
+            raise ValueError(
+                f"gated_delta_core v3 bound {self._g_dtype} log-decay "
+                "and BF16 beta; the host's dtypes moved after binding")
+        if state is None or state.shape != (1, self.heads, 128, 128):
+            raise ValueError("gated_delta_core v3 state shape differs")
+        if state.dtype != self._state_dtype or not state.is_contiguous():
+            raise ValueError(
+                f"gated_delta_core v3 bound contiguous {self._state_dtype} "
+                "state; the host's state moved after binding")
+        # One custom op. The caller's state is read-only and the final state is
+        # written into graph-stable storage for snapshot and rollback.
+        out, state_out = self._step(
+            query[:, 0], key[:, 0], value[:, 0],
+            log_decay[:, 0], beta[:, 0], state,
+            use_qk_l2norm=use_qk_l2norm,
+            state_out=self._state_out,
+            out=self._out,
+        )
+        return out[:, None], state_out if output_final_state else None
+
+
+def bind_gated_delta_core(sample: dict[str, torch.Tensor]):
+    """Bind v3 decode recurrence and launch the observed real sample once.
+
+    The entry is chosen by the observed sample's log-decay dtype — the
+    form the host actually calls with, not a preference."""
+    state = sample.get("state")
+    core = HubV3GatedDeltaCore(
+        sample["query"], g_dtype=sample["g"].dtype,
+        state_dtype=(state.dtype if state is not None
+                     else torch.bfloat16))
+    with torch.no_grad():
+        core(
+            sample["query"], sample["key"], sample["value"],
+            sample["g"], sample["beta"], sample.get("state"),
+            output_final_state=bool(sample.get("output_final_state", True)),
+            use_qk_l2norm=bool(sample.get("use_qk_l2norm", True)),
+        )
+    guard = core._frt_guard
+    if guard is not None:
+        guard.calls = 0
+    return core

flashrt_structures/impls/graph_lowering/__init__.py

@@ -0,0 +1,27 @@
+"""Whole-graph shape-lowering adapters, one per host family.
+
+``capture`` consults this registry when it is handed a model: every
+adapter that recognizes the host pins that family's shape glue for the
+fixed request, and hands back an ``undo``. A host no family recognizes
+is captured as-is — correct for hosts that are already graph-safe.
+"""
+
+from .protocol import (GraphLowering, GraphLoweringRefused,
+                       lower_for_capture,
+                       register_graph_lowering_adapter)
+from .qwen3_vl import Qwen3VLGraphLoweringAdapter
+
+# Built-ins register at import time; they recognize by capability, not
+# by class name or version string.
+register_graph_lowering_adapter(Qwen3VLGraphLoweringAdapter())
+
+from .pi052_denoise import Pi05DenoiseGraphLoweringAdapter  # noqa: E402
+
+register_graph_lowering_adapter(Pi05DenoiseGraphLoweringAdapter())
+
+__all__ = [
+    "GraphLowering",
+    "GraphLoweringRefused",
+    "lower_for_capture",
+    "register_graph_lowering_adapter",
+]

flashrt_structures/impls/graph_lowering/pi052_denoise.py

@@ -0,0 +1,179 @@
+"""Pi0.5 flow-matching family: the timestep schedule becomes resident.
+
+The host builds its denoise schedule inside ``sample_actions`` as
+``torch.tensor([...python floats...], device=cuda)`` — one
+host-to-device copy per call. Whether that line survives capture has
+depended on the compiler's mood: a dynamo that covers the whole method
+bakes it into the graph, a dependency upgrade that adds a graph break
+in front of it drops the copy into the capture stream and the capture
+refuses (measured: the same host line passed on 2026-07-25 and
+refused after a transformers upgrade landed the next day). A pin must
+not gamble on coverage.
+
+The pin scopes one rule around the host's own method: a
+``torch.tensor`` call that constructs a *constant float list* on a
+device resolves to a cached resident tensor — same values, same
+device, same dtype, allocated once outside capture. Everything else
+passes straight through, the schedule stays value-identical by
+construction, and the undo restores the host method bit-for-bit.
+Scalar constructors use device-native fills instead of a CPU staging copy;
+unlike list schedules, mutable scalars are not cached between calls.
+"""
+
+from __future__ import annotations
+
+import types
+
+import torch
+
+from .protocol import GraphLowering
+
+
+def _looks_like_pi05_flow(module) -> bool:
+    return (callable(getattr(module, "sample_actions", None))
+            and callable(getattr(module, "denoise_step", None))
+            and callable(getattr(module, "embed_suffix", None))
+            and hasattr(module, "paligemma_with_expert"))
+
+
+class Pi05DenoiseGraphLoweringAdapter:
+    """Family: pi05_denoise — one pin, the resident step schedule."""
+
+    def lower(self, model, forward) -> GraphLowering | None:
+        target = None
+        if _looks_like_pi05_flow(model):
+            target = model
+        else:
+            for _name, mod in getattr(
+                    model, "named_modules", lambda: ())():
+                if _looks_like_pi05_flow(mod):
+                    target = mod
+                    break
+            if target is None and _looks_like_pi05_flow(
+                    getattr(model, "model", None)):
+                target = model.model
+        if target is None:
+            return None
+
+        cache: dict[tuple, torch.Tensor] = {}
+        real_tensor = torch.tensor
+        host_fn = target.sample_actions
+        had_instance = "sample_actions" in target.__dict__
+
+        def caching_tensor(data, *args, **kwargs):
+            device = kwargs.get("device")
+            if (device is not None and not args
+                    and isinstance(data, (float, int, bool))):
+                return torch.full((), data, **kwargs)
+            if (device is not None and isinstance(data, (list, tuple))
+                    and data
+                    and all(isinstance(x, (float, int, bool))
+                            for x in data)):
+                key = (tuple(data), str(device),
+                       str(kwargs.get("dtype")))
+                hit = cache.get(key)
+                if hit is None:
+                    hit = real_tensor(data, *args, **kwargs)
+                    cache[key] = hit
+                return hit
+            return real_tensor(data, *args, **kwargs)
+
+        real_setitem = torch.Tensor.__setitem__
+
+        def filling_setitem(t, idx, val):
+            # a python-scalar write into a CUDA tensor is a
+            # host-to-device copy the capture stream refuses; the same
+            # store as an immediate-value fill is graph-legal and
+            # bit-identical
+            if (t.is_cuda and isinstance(idx, int)
+                    and isinstance(val, (int, float, bool))):
+                t.narrow(0, idx, 1).fill_(val)
+                return
+            real_setitem(t, idx, val)
+
+        def pinned(self, *args, **kwargs):
+            torch.tensor = caching_tensor
+            torch.Tensor.__setitem__ = filling_setitem
+            try:
+                return host_fn(*args, **kwargs)
+            finally:
+                torch.tensor = real_tensor
+                torch.Tensor.__setitem__ = real_setitem
+
+        target.sample_actions = types.MethodType(pinned, target)
+
+        # ---- pin 3: the pixel-patch embedding stack rides the band ----
+        # The host keeps norms/embeddings in fp32 as a training-fidelity
+        # choice; in the captured serving form every consumer of the
+        # patch embeds casts to bf16 at its own entry, so the fp32 conv
+        # pair is pure spend. Structural match only (a full-patch
+        # Conv2d — kernel == stride — beside a position Embedding),
+        # weights carried down in place with the originals retained,
+        # and the arm's end-to-end parity gate stays the judge.
+        embed_saved: list = []
+        embed_hooks: list = []
+        for _n, mod in getattr(model, "named_modules", lambda: ())():
+            pe = getattr(mod, "patch_embedding", None)
+            pos = getattr(mod, "position_embedding", None)
+            if not (isinstance(pe, torch.nn.Conv2d)
+                    and isinstance(pos, torch.nn.Embedding)):
+                continue
+            if tuple(pe.kernel_size) != tuple(pe.stride):
+                continue
+            f32 = [p for p in mod.parameters()
+                   if p.dtype == torch.float32]
+            if not f32:
+                continue
+            for p in f32:
+                embed_saved.append((p, p.data))
+                p.data = p.data.to(torch.bfloat16)
+            embed_hooks.append(pe.register_forward_pre_hook(
+                lambda _m, args: (args[0].to(torch.bfloat16),)
+                + tuple(args[1:])))
+            # dtype-transparent at the module boundary: every
+            # downstream consumer keeps seeing the dtype the host
+            # chose; only the patch projection itself rides the band
+            embed_hooks.append(mod.register_forward_hook(
+                lambda _m, _a, out: out.to(torch.float32)
+                if isinstance(out, torch.Tensor) else out))
+        # ---- pin 4: fp32 host linears ride the band, transparently --
+        # The same fidelity policy leaves a handful of glue linears
+        # (modality projector, time/action MLPs) in fp32, which on this
+        # class of device means simt kernels with no tensor cores. Each
+        # one is carried down in place with both boundaries cast back,
+        # so every consumer and producer keeps its dtype contract and
+        # the parity gate judges the whole move.
+        for _n, mod in getattr(model, "named_modules", lambda: ())():
+            if not isinstance(mod, torch.nn.Linear):
+                continue
+            if mod.weight.dtype is not torch.float32:
+                continue
+            embed_saved.append((mod.weight, mod.weight.data))
+            mod.weight.data = mod.weight.data.to(torch.bfloat16)
+            if mod.bias is not None:
+                embed_saved.append((mod.bias, mod.bias.data))
+                mod.bias.data = mod.bias.data.to(torch.bfloat16)
+            embed_hooks.append(mod.register_forward_pre_hook(
+                lambda _m, args: (args[0].to(torch.bfloat16),)
+                + tuple(args[1:])))
+            embed_hooks.append(mod.register_forward_hook(
+                lambda _m, _a, out: out.to(torch.float32)))
+        pins = ["resident_step_schedule", "scalar_setitem_fill"]
+        if embed_saved:
+            pins.append("patch_embed_band")
+
+        def undo() -> None:
+            torch.tensor = real_tensor
+            if had_instance:
+                target.sample_actions = host_fn
+            elif "sample_actions" in target.__dict__:
+                del target.sample_actions
+            for hook in embed_hooks:
+                hook.remove()
+            for p, data in embed_saved:
+                p.data = data
+
+        return GraphLowering(
+            undo=undo, family="pi05_denoise",
+            pins=tuple(pins),
+            details={"host": type(target).__name__})

flashrt_structures/impls/graph_lowering/protocol.py

@@ -0,0 +1,76 @@
+"""Host-family adapter protocol for whole-graph shape lowering.
+
+CUDA graph capture needs every shape-derived quantity of the request to
+be a constant: position tables, token routing, sequence cumsums — the
+things a host recomputes per call, often through a synchronize the
+capture cannot record. Which functions those are is host-family
+knowledge, exactly like where the attention math runs — so the lowering
+lives in registered adapters, not in user harnesses. A user asks to
+capture; the family that recognizes its host pins its own glue.
+
+Every pin is a shape-derived constant of one fixed request, never a
+value-dependent quantity, and every application returns an ``undo``
+that restores the host bit-for-bit.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Mapping, Protocol
+
+
+class GraphLoweringRefused(RuntimeError):
+    """A recognized host cannot be pinned safely."""
+
+
+@dataclass(frozen=True)
+class GraphLowering:
+    """One applied family lowering and how to take it back off."""
+
+    undo: Callable[[], None]
+    family: str
+    pins: tuple[str, ...]
+    details: Mapping[str, Any] = field(default_factory=dict)
+
+
+class GraphLoweringAdapter(Protocol):
+    """One host-family realization of the lowering contract."""
+
+    def lower(
+        self,
+        model: Any,
+        forward: Callable[[], Any],
+    ) -> GraphLowering | None: ...
+
+
+_ADAPTERS: list[GraphLoweringAdapter] = []
+
+
+def register_graph_lowering_adapter(adapter: GraphLoweringAdapter) -> None:
+    """Register a host-family lowering adapter."""
+    _ADAPTERS.append(adapter)
+
+
+def lower_for_capture(
+    model: Any,
+    forward: Callable[[], Any],
+) -> list[GraphLowering]:
+    """Apply every registered lowering that recognizes this host.
+
+    An empty list is not a fallback: it means no family recognized the
+    model, and capture proceeds on the host's own forward — which is
+    correct for hosts that are already graph-safe. A recognized family
+    that cannot pin safely raises :class:`GraphLoweringRefused` instead
+    of leaving the host half-pinned; adapters must apply atomically.
+    """
+    applied: list[GraphLowering] = []
+    try:
+        for adapter in _ADAPTERS:
+            lowering = adapter.lower(model, forward)
+            if lowering is not None:
+                applied.append(lowering)
+    except Exception:
+        for lowering in reversed(applied):
+            lowering.undo()
+        raise
+    return applied