flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/norm_fused/bf16.py

@@ -0,0 +1,94 @@
+"""norm_fused — a plain norm run by a fused kernel at compute dtype.
+
+Vision towers commonly keep their LayerNorms in FP32 while the rest of
+the block runs in BF16, so every norm pays a dtype round trip plus an
+unfused mean/variance pass. This replacement runs the norm in one fused
+BF16 kernel and hands the result back in the host's dtype.
+
+Qualification is the host's own dtype: the win comes from collapsing an
+FP32 norm into a fused BF16 one, so a norm the host already runs in
+BF16 is left alone (there is nothing to collapse), and the parity gate
+adjudicates the numerical difference the dtype change introduces.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/flashrt-residual-norm-quant",
+    "version": ">=1",
+}
+
+
+class FusedNorm(GuardedSeam, torch.nn.Module):
+    """Drop-in for an affine LayerNorm, computed by a fused kernel."""
+
+    _frt_host_attr = "host_norm"
+    _frt_can_fallback = True
+
+    def __init__(self, original: torch.nn.Module):
+        super().__init__()
+        self.host_norm = original
+        ks = hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+        self._fn = ks.layer_norm_bf16
+        self.register_buffer("w", original.weight.detach().to(
+            torch.bfloat16))
+        self.register_buffer("b", original.bias.detach().to(
+            torch.bfloat16))
+        self.eps = float(getattr(original, "eps", 1e-6))
+        self._frt_arm(dtypes=CAST_OK, device=self.w.device,
+                      k=int(self.w.shape[0]))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        # the kernel's contract is 2D [rows, width]; hosts hand the norm
+        # whatever leading shape their block carries
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous()
+        y = self._fn(flat, self.w, self.b, self.eps)
+        return y.reshape(shape).to(x.dtype)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(super().__getattr__("host_norm"), name)
+
+
+def bind_norm_fused(original: torch.nn.Module,
+                    host_dtypes=None) -> FusedNorm:
+    """Bind a fused norm, refusing where there is nothing to collapse.
+
+    ``host_dtypes`` is the set of input dtypes this norm was observed with
+    during calibration — one observation, not a statistic. A host already
+    running the norm at a compute dtype has nothing for this structure to
+    collapse, and the refusal names the dtype so it reads as "not in this
+    form" rather than "not supported".
+    """
+    if getattr(original, "weight", None) is None or \
+            getattr(original, "bias", None) is None:
+        raise ValueError("norm_fused: needs an affine norm (weight+bias)")
+    if host_dtypes and torch.float32 not in set(host_dtypes):
+        raise ValueError(
+            "norm_fused: host already runs this norm at compute "
+            f"dtype ({sorted(str(d) for d in host_dtypes)}) — nothing "
+            "to collapse")
+    bound = FusedNorm(original)
+    # bind-time smoke through the real entry point, at a 3D host shape:
+    # a stale build, a missing symbol, or a kernel whose rank contract
+    # moved must surface here as a clean bind refusal, not mid-forward
+    probe_in = torch.zeros(1, 2, bound.w.shape[0], device=bound.w.device)
+    probe = bound(probe_in)
+    if probe.shape != probe_in.shape or not torch.isfinite(probe).all():
+        raise ValueError(
+            f"refused: norm_fused bind smoke produced shape "
+            f"{tuple(probe.shape)}, "
+            f"finite={bool(torch.isfinite(probe).all())}")
+    return bound

flashrt_structures/impls/norm_fused/fp8_producer.py

@@ -0,0 +1,84 @@
+"""An affine LayerNorm that emits FP8 directly, at a consumer's scale.
+
+The pipeline fact this serves: a vision block's pre-FFN norm output has
+exactly one consumer, the FFN — and when that FFN is seated in FP8 form
+its first act is to quantize its input. Emitting FP8 from the norm
+itself (one fused kernel: no-affine LN + scale/shift + static quantize,
+with scale=(gamma-1), shift=beta reproducing the affine norm exactly)
+deletes the FFN's own input quantize and the norm's dtype round trip.
+
+The hard precedent this respects: handing FP8 to a *host* consumer is
+garbage-in-silence (measured 0.24 output match at the decoder norm
+boundary). This producer is therefore only ever seated by the
+negotiation pass that pairs it with an FP8-input seat as the direct
+consumer — seat produces, seat consumes, and the consumer's FP8_ONLY
+guard refuses loudly if anything else arrives between them. Whether
+the pair actually pays is measured at bind, never assumed.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/adaptive-layernorm-producers",
+    "version": ">=1",
+}
+
+
+class FusedNormFp8Producer(GuardedSeam, torch.nn.Module):
+    """Drop-in for an affine LayerNorm whose sole consumer eats FP8."""
+
+    _frt_host_attr = "host_norm"
+    _frt_can_fallback = False   # the consumer expects FP8; a BF16
+    # fallback here would feed the paired seat out of contract, so an
+    # out-of-form input must refuse loudly instead of degrading quietly
+
+    def __init__(self, original: torch.nn.Module,
+                 act_scale: torch.Tensor):
+        super().__init__()
+        self.host_norm = original
+        ks = hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+        self._fn = ks.ada_layer_norm_quant_fp8_bf16
+        # scale=(gamma-1), shift=beta: norm(x)*(1+scale)+shift is then
+        # exactly the host's affine LayerNorm, quantized
+        self.register_buffer("w", (original.weight.detach().float() - 1.0)
+                             .to(torch.bfloat16).contiguous())
+        self.register_buffer("b", original.bias.detach()
+                             .to(torch.bfloat16).contiguous())
+        self.register_buffer("act_scale",
+                             act_scale.detach().reshape(1).float())
+        self.eps = float(getattr(original, "eps", 1e-6))
+        self._frt_arm(dtypes=CAST_OK, device=self.w.device,
+                      k=int(self.w.shape[0]))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        shape = x.shape
+        flat = x.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous()
+        y = self._fn(flat, self.w, self.b, self.act_scale, self.eps)
+        return y.reshape(shape)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_norm":
+                raise
+            return getattr(super().__getattr__("host_norm"), name)
+
+
+def bind_norm_fp8_producer(original: torch.nn.Module,
+                           act_scale: torch.Tensor
+                           ) -> FusedNormFp8Producer:
+    if getattr(original, "weight", None) is None \
+            or getattr(original, "bias", None) is None:
+        raise ValueError("fp8 norm producer needs a two-sided affine "
+                         "LayerNorm")
+    return FusedNormFp8Producer(original, act_scale)

flashrt_structures/impls/patch_projection/__init__.py

@@ -0,0 +1,3 @@
+from .bf16_flat import FlatPatchProjection, bind_flat_patch_projection
+
+__all__ = ["FlatPatchProjection", "bind_flat_patch_projection"]

flashrt_structures/impls/patch_projection/bf16_flat.py

@@ -0,0 +1,139 @@
+"""BF16 lowering for processor-preflattened full-patch Conv3D modules."""
+
+from __future__ import annotations
+
+from typing import Mapping, Sequence
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/flashrt-gemm-epilogues",
+    "version": ">=1",
+}
+
+
+def _kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+
+
+class FlatPatchProjection(GuardedSeam, torch.nn.Module):
+    """Drop-in replacement for an exact full-patch Conv3D wrapper.
+
+    The host owns checkpoint-layout Conv3D weights ``[N,C,T,P,P]``. Binding
+    flattens them to the Hub API's ``[K,N]`` GEMM layout exactly once. The
+    retained host module remains the fallback and state-dict owner.
+    """
+
+    _frt_host_attr = "host_patch"
+    _frt_can_fallback = True
+
+    def __init__(
+        self,
+        weight_kn: torch.Tensor,
+        bias: torch.Tensor | None,
+        *,
+        row_capacity: int,
+        host_dtypes: Sequence[torch.dtype],
+        original: torch.nn.Module,
+        kernel=None,
+    ) -> None:
+        super().__init__()
+        self._weight_kn = weight_kn
+        self._bias = bias
+        self._row_capacity = int(row_capacity)
+        self.host_patch = original
+        self._ops = _kernel() if kernel is None else kernel
+        entry = (
+            "bf16_linear_bias_bf16"
+            if bias is not None
+            else "bf16_linear_bf16"
+        )
+        try:
+            self._fn = getattr(self._ops, entry)
+        except AttributeError as exc:
+            raise ValueError(
+                f"patch_projection Hub artifact lacks {entry}"
+            ) from exc
+        self._out = torch.empty(
+            self._row_capacity,
+            weight_kn.shape[1],
+            device=weight_kn.device,
+            dtype=torch.bfloat16,
+        )
+        self._frt_arm(
+            dtypes=tuple(host_dtypes) or CAST_OK,
+            device=weight_kn.device,
+            k=int(weight_kn.shape[0]),
+            row_capacity=self._row_capacity,
+        )
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_patch":
+                raise
+            return getattr(super().__getattr__("host_patch"), name)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        flat = x.reshape(-1, self._weight_kn.shape[0])
+        rows = flat.shape[0]
+        out = self._out[:rows]
+        flat = flat.to(torch.bfloat16).contiguous()
+        if self._bias is None:
+            return self._fn(flat, self._weight_kn, out=out)
+        return self._fn(flat, self._weight_kn, self._bias, out=out)
+
+
+@torch.no_grad()
+def bind_flat_patch_projection(
+    weights: Mapping[str, torch.Tensor],
+    *,
+    row_profile: Sequence[int],
+    host_dtypes: Sequence[torch.dtype],
+    original: torch.nn.Module,
+) -> FlatPatchProjection:
+    """Bind a full-patch projection from checkpoint weights ``w[N,K]``."""
+    if not row_profile:
+        raise ValueError("patch_projection: no real patch rows were observed")
+    w = weights["w"]
+    if w.dim() != 2 or w.dtype is not torch.bfloat16 or not w.is_cuda:
+        raise ValueError(
+            "patch_projection requires CUDA BF16 checkpoint weights [N,K]"
+        )
+    b = weights.get("b")
+    if b is not None:
+        if b.shape != (w.shape[0],):
+            raise ValueError("patch_projection bias width does not match N")
+        b = b.detach().to(torch.bfloat16).contiguous()
+    weight_kn = w.detach().t().contiguous()
+    capacity = max(int(row) for row in row_profile)
+    bound = FlatPatchProjection(
+        weight_kn,
+        b,
+        row_capacity=capacity,
+        host_dtypes=host_dtypes,
+        original=original,
+    )
+    # A fallback-capable seam must prove the formal artifact launches at bind
+    # time; otherwise a stale package would look numerically perfect by
+    # silently running the retained host.
+    sample_dtype = next(iter(host_dtypes), torch.bfloat16)
+    sample = torch.zeros(
+        capacity,
+        w.shape[1],
+        device=w.device,
+        dtype=sample_dtype,
+    )
+    bound(sample)
+    if bound._frt_guard is not None:
+        bound._frt_guard.calls = 0
+    return bound

flashrt_structures/impls/prefill_tower/__init__.py

@@ -0,0 +1,10 @@
+"""The plain-norm decoder tower region: the prefix-building pass.
+
+The tower that fills the cache another stack later attends over: plain
+RMS norms (no conditioning), bias-free attention and gated FFN, one
+forward per observation with ``use_cache``. Its chain candidate
+re-expresses the per-layer loop in static-FP8 hub primitives while
+writing host-layout keys back into the host's own cache, so every
+downstream consumer — the sibling chain or the host fallback — reads
+what it always read.
+"""