flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/decoder_block/__init__.py

@@ -0,0 +1,6 @@
+from .attn_sublayer import (PackedAttnSublayer, bind_attn_sublayer,
+                            rotate_half)
+from .fused import FusedDecoderBlock, bind_decoder_block, qualify
+
+__all__ = ["FusedDecoderBlock", "PackedAttnSublayer", "bind_attn_sublayer",
+           "bind_decoder_block", "qualify", "rotate_half"]

flashrt_structures/impls/decoder_block/attn_sublayer.py

@@ -0,0 +1,110 @@
+"""The attention sublayer with no layout churn between its parts.
+
+Three bound structures already sit inside a host attention module: the
+packed projections, the rotary embedding, the fused attention core. Each
+is faster than what it replaced, and between them the host still pays for
+a layout it does not need. The host lays q/k/v out as ``(B, H, S, D)``
+because that is what eager SDPA wants; the fused kernel wants
+``(B, S, H, D)``, which is exactly what the projections' own output view
+already is. So the host transposes, the rotary embedding runs on the
+transposed layout, and the core transposes back and makes it
+contiguous — two cancelling transposes plus the copies around them, per
+projection, per layer, per step.
+
+None of the three seams can see that, because each is bound inside the
+module that owns the layout. The sublayer boundary can: run the packed
+projections, view their output as ``(B, S, H, D)``, apply the rotary
+embedding on that layout (the ``unsqueeze`` axis moves, the arithmetic
+does not), and hand it straight to the kernel.
+
+The rotary form is family-specific (half-split against interleaved), so
+it enters as a callable from the host-family adapter; everything else
+here is the generic pre-norm attention sublayer.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Half-split rotation (Llama/Gemma/Qwen convention)."""
+    half = x.shape[-1] // 2
+    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+
+class PackedAttnSublayer(nn.Module):
+    """Packed projections -> rotary -> fused core -> output projection."""
+
+    def __init__(self, attn: nn.Module, core, *, scale: float,
+                 rotate=rotate_half, q_heads: int = 0):
+        super().__init__()
+        self.attn = attn
+        self.core = core
+        self.scale = scale
+        self.rotate = rotate
+        self.q_heads = q_heads
+
+    def forward(self, x: torch.Tensor, position_embeddings=None, **kw):
+        a = self.attn
+        bsz, seq, _ = x.shape
+        hd = a.head_dim
+        pack = a.q_proj if getattr(a.q_proj, "joint_slots", 0) else None
+        if pack is not None:
+            # q and k are one contiguous run of the packed output and
+            # the rotary embedding is the same arithmetic on both, so it
+            # runs once over the pair. Splitting them first would cost a
+            # kernel and the copies that separate them.
+            qk = pack.joint(x).view(bsz, seq, -1, hd)
+            qk = self._rope(qk, position_embeddings)
+            q, k = qk[:, :, :self.q_heads], qk[:, :, self.q_heads:]
+            v = a.v_proj(x).view(bsz, seq, -1, hd)
+        else:
+            # the host's own call order is the data dependency the
+            # packed projection relies on: the first call runs the GEMM,
+            # the others read its stash
+            q = a.q_proj(x).view(bsz, seq, -1, hd)
+            k = a.k_proj(x).view(bsz, seq, -1, hd)
+            v = a.v_proj(x).view(bsz, seq, -1, hd)
+            q = self._rope(q, position_embeddings)
+            k = self._rope(k, position_embeddings)
+        out = self.core.forward_suffix(q, k, v, scale=self.scale)
+        return a.o_proj(out.reshape(bsz, seq, -1))
+
+    def _rope(self, t: torch.Tensor, position_embeddings):
+        if position_embeddings is None:
+            return t
+        cos, sin = position_embeddings
+        cos = cos.unsqueeze(2).to(t.dtype)
+        sin = sin.unsqueeze(2).to(t.dtype)
+        return t * cos + self.rotate(t) * sin
+
+
+def bind_attn_sublayer(attn: nn.Module, core, *, rotate=rotate_half):
+    """Compose one attention sublayer around an already-bound core.
+
+    Returns ``None`` rather than raising when the host module or the core
+    is missing a part: the block then keeps the host's own attention, so
+    this can only add coverage.
+    """
+    if core is None or not hasattr(core, "forward_suffix"):
+        return None
+    for attr in ("q_proj", "k_proj", "v_proj", "o_proj", "head_dim"):
+        if not hasattr(attn, attr):
+            return None
+    # the sublayer's projections produce the new tokens only, so the
+    # packed plan's suffix has to be exactly those. A host that carries
+    # its own KV cache into attention breaks that equality, and there the
+    # host's attention module stays.
+    plan = getattr(core, "plan", None)
+    if plan is None or plan.suffix_len != getattr(core, "seq_q", -1):
+        return None
+    scale = getattr(attn, "scaling", None)
+    if scale is None:
+        scale = getattr(attn, "scale", None)
+    if scale is None:
+        return None
+    q_heads = attn.q_proj.out_features // attn.head_dim
+    return PackedAttnSublayer(attn, core, scale=float(scale),
+                              rotate=rotate, q_heads=q_heads)

flashrt_structures/impls/decoder_block/fused.py

@@ -0,0 +1,167 @@
+"""decoder_block — the pre-norm transformer block as one boundary.
+
+Every region structure in this library binds inside a block: the norm,
+the packed projections, the attention, the MLP. What none of them can
+see is the *dataflow between them* — and that is where the residual
+adds, the gate broadcasts and the dtype round trips live. A host block
+is a fixed shape:
+
+    r = h;  h, g = norm_in(h, cond);   h = attn(h);  h = r + h * g
+    r = h;  h, g = norm_out(h, cond);  h = mlp(h);   h = r + h * g
+
+so the second norm's real input is not a hidden state, it is a pending
+``residual + attn_out * gate``. The adaptive-norm kernel already takes
+exactly that (it computes the gated residual, norms it, modulates and
+quantizes in one pass) — bound at the norm boundary there is nothing to
+hand it, so the residual argument gets zeros, the host keeps its own
+elementwise add, and the fused producer looks like dead weight. It was
+measured as such and refused. The boundary was wrong, not the kernel.
+
+Binding the block puts the residual back in the producer's hands. Three
+things follow, none of which is available one seam at a time:
+
+- the pending residual add disappears into the producer's kernel;
+- the producer emits FP8, so the MLP takes the FP8 entry and its own
+  input quantization goes away;
+- the step lookup is resolved once and shared by both producers instead
+  of being recomputed per norm.
+
+The block owns composition, not kernels: every slot is a module bound by
+its own structure, and a slot that did not bind keeps the host's child.
+It can therefore only add to what the region structures already do.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+_BLOCK_ATTRS = ("self_attn", "mlp", "input_layernorm",
+                "post_attention_layernorm")
+
+
+def _parameterised_children(module: nn.Module) -> set[str]:
+    return {name for name, child in module.named_children()
+            if any(True for _ in child.parameters())}
+
+
+class FusedDecoderBlock(GuardedSeam, nn.Module):
+    """Pre-norm block whose sublayer dataflow runs through the producers.
+
+    The host block is retained whole, which makes this the widest way back
+    in the library: a call outside the calibrated form runs the entire
+    original block. That is exact rather than approximate, because the
+    block never swapped the host's own children — it holds separately
+    bound copies and leaves the host's sublayers where they were.
+    """
+
+    _frt_host_attr = "host"
+    _frt_can_fallback = True
+
+    def __init__(self, host: nn.Module, producer_in, producer_out,
+                 ffn: nn.Module, *, cond_kw: str = "adarms_cond",
+                 returns_tuple: bool = False, attn=None):
+        super().__init__()
+        self.host = host
+        self.producer_in = producer_in
+        self.producer_out = producer_out
+        self.ffn = ffn
+        self.cond_kw = cond_kw
+        self.returns_tuple = returns_tuple
+        # an attention sublayer that owns its own layout, or None to keep
+        # the host's attention module (which owns the layout itself)
+        self.attn = attn
+        self.own_attn = attn is not None
+        rows, dim = producer_in.resid.shape
+        self._frt_arm(dtypes=CAST_OK, device=producer_in.resid.device,
+                      k=int(dim), rows=int(rows))
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
+        # before the conditioning is taken out of kwargs: the host block
+        # expects its own signature back if this call has to go to it
+        admitted = self._frt_admit(hidden_states, *args, **kwargs)
+        if admitted is not PROCEED:
+            return admitted
+        cond = kwargs.pop(self.cond_kw, None)
+        idx = self.producer_in.resolve(cond)
+
+        y, gate = self.producer_in.produce(hidden_states, idx)
+        y = y.reshape(hidden_states.shape)
+        if self.own_attn:
+            attn_out = self.attn(
+                y, position_embeddings=kwargs.get("position_embeddings"))
+        else:
+            attn_out = self.host.self_attn(y, *args, **kwargs)
+            if isinstance(attn_out, tuple):
+                attn_out = attn_out[0]
+
+        # the pending "hidden_states + attn_out * gate" is the second
+        # norm's input; the kernel takes it whole
+        resid, y, gate = self.producer_out.absorb(
+            hidden_states, attn_out, gate, idx)
+
+        out = resid + self.ffn(y) * gate
+        out = out.reshape(hidden_states.shape)
+        return (out,) if self.returns_tuple else out
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(super().__getattr__("host"), name)
+
+
+def qualify(host: nn.Module) -> None:
+    """Refuse a host block whose dataflow this structure does not model.
+
+    The check is structural, not behavioural: the four sublayer slots
+    must be there and nothing else that carries weights may be, because
+    a fifth parameterised child (a second pair of norms, a cross
+    attention) is a sublayer this block would silently drop. What the
+    forward then does with those slots is adjudicated by the parity gate
+    at the tick boundary, which is where this structure's reference is
+    declared.
+    """
+    missing = [a for a in _BLOCK_ATTRS
+               if not isinstance(getattr(host, a, None), nn.Module)]
+    if missing:
+        raise ValueError(f"decoder_block: host lacks {missing}")
+    extra = _parameterised_children(host) - set(_BLOCK_ATTRS)
+    if extra:
+        raise ValueError(
+            "decoder_block: host carries sublayers this structure does "
+            f"not model ({sorted(extra)}) — keeping the host block")
+
+
+def bind_decoder_block(host: nn.Module, producer_in, producer_out,
+                       ffn: nn.Module, *, cond_kw: str = "adarms_cond",
+                       returns_tuple: bool = False,
+                       attn=None) -> FusedDecoderBlock:
+    """Compose bound sublayer structures into one block.
+
+    ``producer_in`` / ``producer_out`` are bound ``adaln_producer``
+    modules on the same conditioning stream (they must share a locator
+    for the step lookup to be shared); ``ffn`` is a bound
+    ``decoder_ffn`` on the FP8 entry, which is what makes the producer's
+    quantize load-bearing.
+    """
+    qualify(host)
+    for name, prod in (("producer_in", producer_in),
+                       ("producer_out", producer_out)):
+        if not hasattr(prod, "resolve"):
+            raise ValueError(
+                f"decoder_block: {name} is not an adaln_producer")
+    if not producer_out.can_absorb:
+        raise ValueError(
+            "decoder_block: the second producer cannot absorb a residual "
+            "(it is not the rms form with fp8 output) — without that "
+            "fold the block boundary buys nothing over the region seams")
+    if producer_in.locator is not producer_out.locator:
+        raise ValueError(
+            "decoder_block: the two producers do not share a step "
+            "locator, so the lookup would still run twice")
+    return FusedDecoderBlock(host, producer_in, producer_out, ffn,
+                             cond_kw=cond_kw,
+                             returns_tuple=returns_tuple, attn=attn)

flashrt_structures/impls/decoder_ffn/fp8_static.py

@@ -0,0 +1,310 @@
+"""FP8-static implementation of the ``decoder_ffn`` structure.
+
+Composes the fused FP8 gate/up -> activation -> down block from the
+``flashrt/flashrt-fp8-swiglu-ffn`` Hub kernel behind the structure
+boundary. Two bind entrypoints share the packing and calibration code:
+``bind`` covers the full structure (norm and AdaLN modulation run in
+torch ahead of the fused block); ``bind_mlp_seam`` covers the
+normed-input -> ffn-output slice for hosts whose replaceable module
+boundary is the MLP. Activation scales are static per-tensor,
+calibrated from caller-provided representative inputs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Callable, Mapping, Sequence
+
+import torch
+
+from ...guard import CAST_OK, FP8_ONLY, PROCEED, GuardedSeam
+
+KERNEL_DEP = {
+    "provider": "hf",
+    "repo": "flashrt/flashrt-fp8-swiglu-ffn",
+    "version": ">=1",
+}
+
+_FP8 = torch.float8_e4m3fn
+_FP8_MAX = 448.0
+_ENTRYPOINTS = {"gelu": "bf16_fp8_geglu_mlp_bf16",
+                "silu": "bf16_fp8_swiglu_mlp_bf16"}
+# fp8 entry: the upstream producer already emitted fp8 with the shared
+# activation scale, so the kernel's own input quantization is dead work.
+# Same math, one less kernel per call.
+_ENTRYPOINTS_FP8 = {"gelu": "fp8_geglu_mlp_bf16",
+                    "silu": "fp8_swiglu_mlp_bf16"}
+
+SUPPORT = {
+    "D": {"min": 512, "max": 16384},
+    "F": {"min": 1024, "max": 16384},
+    "m_classes": ("micro", "small", "medium"),
+}
+
+
+@lru_cache(maxsize=1)
+def _kernel():
+    from flashrt_structures.impls import hub_kernel
+
+    return hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+
+
+def _activation(variant: Mapping[str, str]) -> tuple[str, Callable]:
+    name = variant.get("activation", "gelu")
+    if name not in _ENTRYPOINTS:
+        raise ValueError(f"unsupported activation: {name!r}")
+    if name == "gelu":
+        return name, lambda t: torch.nn.functional.gelu(t, approximate="tanh")
+    return name, torch.nn.functional.silu
+
+
+def _amax_scale(tensor: torch.Tensor) -> torch.Tensor:
+    return (tensor.float().abs().max() / _FP8_MAX).clamp(min=1e-8)
+
+
+def _quantize(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return (tensor.float() / scale).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8)
+
+
+def _normalize(
+    x: torch.Tensor,
+    w_norm: torch.Tensor,
+    mode: str,
+    cond_scale: torch.Tensor | None,
+    cond_shift: torch.Tensor | None,
+    eps: float,
+) -> torch.Tensor:
+    h = x.float()
+    h = h * torch.rsqrt(h.pow(2).mean(dim=-1, keepdim=True) + eps)
+    if mode == "offset":
+        h = h * (1.0 + w_norm.float())
+    elif mode == "direct":
+        h = h * w_norm.float()
+    else:
+        raise ValueError(f"unknown norm_weight_mode: {mode!r}")
+    if cond_scale is not None:
+        h = h * (1.0 + cond_scale.float())
+    if cond_shift is not None:
+        h = h + cond_shift.float()
+    return h.to(torch.bfloat16)
+
+
+def _check_and_pack(weights: Mapping[str, torch.Tensor]):
+    """Validate dims against the support envelope; pack FP8 weights."""
+    w_gate, w_up, w_down = weights["w_gate"], weights["w_up"], weights["w_down"]
+    dim_d, dim_f = w_gate.shape
+    if w_up.shape != (dim_d, dim_f) or w_down.shape != (dim_f, dim_d):
+        raise ValueError(
+            f"inconsistent weight dims: gate {tuple(w_gate.shape)}, "
+            f"up {tuple(w_up.shape)}, down {tuple(w_down.shape)}"
+        )
+    for name, dim in (("D", dim_d), ("F", dim_f)):
+        bounds = SUPPORT[name]
+        if not bounds["min"] <= dim <= bounds["max"]:
+            raise ValueError(
+                f"{name}={dim} outside support envelope "
+                f"[{bounds['min']}, {bounds['max']}]"
+            )
+    if not (w_gate.is_cuda and w_up.is_cuda and w_down.is_cuda):
+        raise ValueError("fp8_static requires CUDA-resident weights")
+    gate_up = torch.cat([w_gate.t(), w_up.t()], dim=0).contiguous()
+    down = w_down.t().contiguous()
+    return gate_up, down, _amax_scale(gate_up), _amax_scale(down)
+
+
+def _calibrate_scales(
+    normed_samples: Sequence[torch.Tensor],
+    w_gate: torch.Tensor,
+    w_up: torch.Tensor,
+    act: Callable[[torch.Tensor], torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Static per-tensor input/hidden scales from normed activations."""
+    if not normed_samples:
+        raise ValueError("calibration samples must be non-empty")
+    device = w_gate.device
+    input_amax = torch.zeros((), device=device)
+    hidden_amax = torch.zeros((), device=device)
+    with torch.no_grad():
+        for h in normed_samples:
+            flat = h.reshape(-1, h.shape[-1]).float().to(device)
+            hidden = act(flat @ w_gate.float()) * (flat @ w_up.float())
+            input_amax = torch.maximum(input_amax, flat.abs().max())
+            hidden_amax = torch.maximum(hidden_amax, hidden.abs().max())
+    return ((input_amax / _FP8_MAX).clamp(min=1e-8),
+            (hidden_amax / _FP8_MAX).clamp(min=1e-8))
+
+
+@dataclass(frozen=True)
+class BoundDecoderFfnFp8:
+    """Bound callable for the full structure boundary."""
+
+    fused_mlp: Callable[..., torch.Tensor]
+    w_norm: torch.Tensor
+    gate_up_fp8: torch.Tensor
+    down_fp8: torch.Tensor
+    input_scale: torch.Tensor
+    gate_up_scale: torch.Tensor
+    hidden_scale: torch.Tensor
+    down_scale: torch.Tensor
+    norm_weight_mode: str
+    eps: float
+    in_dtype: str = "bf16"
+
+    def ffn(self, normed: torch.Tensor) -> torch.Tensor:
+        """The normed-input -> ffn-output slice (no norm, no residual).
+
+        On the BF16 entry the kernel quantizes the input itself; on the
+        FP8 entry the producer already did, and the input passes
+        straight through."""
+        shape = normed.shape
+        if getattr(self, "in_dtype", "bf16") == "fp8_static":
+            out = self.fused_mlp(
+                normed.reshape(-1, shape[-1]),
+                self.gate_up_fp8, self.down_fp8,
+                self.input_scale.view(1), self.gate_up_scale.view(1),
+                self.hidden_scale.view(1), self.down_scale.view(1))
+            return out.reshape(*shape[:-1], out.shape[-1])
+        out = self.fused_mlp(
+            normed.reshape(-1, shape[-1]).to(torch.bfloat16).contiguous(),
+            self.gate_up_fp8,
+            self.down_fp8,
+            self.input_scale.view(1),
+            self.gate_up_scale.view(1),
+            self.hidden_scale.view(1),
+            self.down_scale.view(1),
+        )
+        return out.reshape(shape).to(normed.dtype)
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        *,
+        cond_scale: torch.Tensor | None = None,
+        cond_shift: torch.Tensor | None = None,
+        cond_gate: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        h = _normalize(x, self.w_norm, self.norm_weight_mode,
+                       cond_scale, cond_shift, self.eps)
+        out = self.ffn(h)
+        if cond_gate is not None:
+            out = out * cond_gate
+        return x + out.to(x.dtype)
+
+
+class FusedGeGluMlp(GuardedSeam, torch.nn.Module):
+    """MLP-seam module for hosts whose replaceable boundary is the MLP.
+
+    The host keeps its own norm, AdaLN gate, and residual. ``original``
+    is retained whole (host MLP naming varies across model families), and
+    attribute lookups fall through to it so hosts that introspect the
+    projection attributes of the module they call keep working. Retaining
+    it is also what makes the seam reversible per call: an input outside
+    the calibrated form runs the host MLP instead of this kernel.
+    """
+
+    _frt_host_attr = "host_mlp"
+    _frt_can_fallback = True
+
+    def __init__(self, bound: BoundDecoderFfnFp8,
+                 original: torch.nn.Module | None = None):
+        super().__init__()
+        self._bound = bound
+        if original is not None:
+            self.host_mlp = original
+        self._frt_arm(
+            dtypes=(FP8_ONLY if bound.in_dtype == "fp8_static" else CAST_OK),
+            device=bound.gate_up_fp8.device,
+            k=int(bound.gate_up_fp8.shape[1]))
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            if name == "host_mlp":
+                raise
+            return getattr(super().__getattr__("host_mlp"), name)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(hidden)
+        if admitted is not PROCEED:
+            return admitted
+        return self._bound.ffn(hidden)
+
+
+def _build(weights, variant, input_scale, hidden_scale, eps):
+    name, _ = _activation(variant)
+    gate_up, down, gate_up_scale, down_scale = _check_and_pack(weights)
+    in_dtype = variant.get("in_dtype", "bf16")
+    table = (_ENTRYPOINTS_FP8 if in_dtype == "fp8_static"
+             else _ENTRYPOINTS)
+    return BoundDecoderFfnFp8(
+        fused_mlp=getattr(_kernel(), table[name]),
+        w_norm=weights["w_norm"],
+        gate_up_fp8=_quantize(gate_up, gate_up_scale),
+        down_fp8=_quantize(down, down_scale),
+        input_scale=input_scale,
+        gate_up_scale=gate_up_scale,
+        hidden_scale=hidden_scale,
+        down_scale=down_scale,
+        in_dtype=in_dtype,
+        norm_weight_mode=variant.get("norm_weight_mode", "offset"),
+        eps=eps,
+    )
+
+
+@torch.no_grad()
+def bind(
+    weights: Mapping[str, torch.Tensor],
+    *,
+    variant: Mapping[str, str],
+    calibration_inputs: Sequence[Mapping[str, torch.Tensor]],
+    eps: float = 1e-6,
+) -> BoundDecoderFfnFp8:
+    """Bind the full structure: calibration inputs are boundary inputs.
+
+    ``calibration_inputs`` must be drawn from the real input distribution
+    of the target binding; static FP8 scales are only as trustworthy as
+    the data they were measured on.
+    """
+    if not calibration_inputs:
+        raise ValueError("calibration_inputs must be non-empty")
+    _, act = _activation(variant)
+    mode = variant.get("norm_weight_mode", "offset")
+    normed = [
+        _normalize(sample["x"], weights["w_norm"], mode,
+                   sample.get("cond_scale"), sample.get("cond_shift"), eps)
+        for sample in calibration_inputs
+    ]
+    input_scale, hidden_scale = _calibrate_scales(
+        normed, weights["w_gate"], weights["w_up"], act)
+    return _build(weights, variant, input_scale, hidden_scale, eps)
+
+
+@torch.no_grad()
+def bind_mlp_seam(
+    weights: Mapping[str, torch.Tensor],
+    *,
+    variant: Mapping[str, str],
+    input_scale: float,
+    hidden_scale: float,
+    original: torch.nn.Module | None = None,
+    eps: float = 1e-6,
+) -> FusedGeGluMlp:
+    """Bind the MLP-seam slice from two already-calibrated scales.
+
+    The scales arrive measured, not derived: ``input_scale`` is the amax at
+    this MLP's input and ``hidden_scale`` the amax at its down
+    projection's input — which is exactly the gated activation this kernel
+    quantises. Recomputing the second one here would mean keeping the
+    seam's inputs alive to run gate/up over them again, and the amax it
+    would arrive at is the one the host already produced.
+
+    Both are per-tensor FP8 scales (amax/448), reduced across calibration
+    samples by the caller through ``flash_rt.core.calibration``.
+    """
+    dev = weights["w_gate"].device
+    bound = _build(weights, variant,
+                   torch.tensor(float(input_scale), device=dev),
+                   torch.tensor(float(hidden_scale), device=dev), eps)
+    return FusedGeGluMlp(bound, original=original)

flashrt_structures/impls/decoder_ffn/fp8_static.yaml

@@ -0,0 +1,22 @@
+impl: fp8_static
+structure: decoder_ffn
+version: 1
+recipe: fp8_static
+backends: [cuda, rocm]
+form: hub
+module: fp8_static
+entrypoint: bind
+
+kernels:
+  - {provider: hf, repo: flashrt/flashrt-fp8-swiglu-ffn, version: ">=1"}
+
+envelope:
+  support:
+    D: {min: 512, max: 16384}
+    F: {min: 1024, max: 16384}
+    align: {}
+    m_classes: [micro, small, medium]
+  notes: >
+    Support bounds mirror the checks enforced in bind(); a kernel-side
+    can_implement probe should replace them once exposed. Activation
+    scales are static per-tensor and require calibration inputs at bind.