flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/prefill_tower/region.py

@@ -0,0 +1,99 @@
+"""The prefill_tower region family: structural identification, candidates.
+
+The identifier matches shape, never names: a module carrying a
+``layers`` list whose blocks each hold a bias-free attention group and
+gated FFN under *plain* affine RMS norms (a 1-D weight and no
+conditioning projection — the conditioned twin belongs to the sibling
+family), plus a rotary table and a plain stack-level norm. The key
+width must be narrower than the query width: the chain's cache layout
+is written for the single-KV band.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from . import fp8_chain
+from ... import regions
+
+
+def _plain_linear(mod) -> bool:
+    return isinstance(mod, torch.nn.Linear) and mod.bias is None
+
+
+def _plain_norm(mod) -> bool:
+    return fp8_chain._plain_norm_weight(mod) is not None
+
+
+def _block_ok(block) -> bool:
+    attn = getattr(block, "self_attn", None)
+    mlp = getattr(block, "mlp", None)
+    if attn is None or mlp is None:
+        return False
+    if not all(_plain_linear(getattr(attn, a, None))
+               for a in ("q_proj", "k_proj", "v_proj", "o_proj")):
+        return False
+    if not all(_plain_linear(getattr(mlp, a, None))
+               for a in ("gate_proj", "up_proj", "down_proj")):
+        return False
+    if mlp.gate_proj.out_features != mlp.up_proj.out_features:
+        return False
+    if not _plain_norm(getattr(block, "input_layernorm", None)):
+        return False
+    return _plain_norm(getattr(block, "post_attention_layernorm", None))
+
+
+def identify(model) -> list[str]:
+    roots = []
+    for path, mod in model.named_modules():
+        layers = getattr(mod, "layers", None)
+        if not isinstance(layers, torch.nn.ModuleList) or len(layers) < 2:
+            continue
+        if not callable(getattr(mod, "rotary_emb", None)):
+            continue
+        head = layers[0]
+        attn = getattr(head, "self_attn", None)
+        if attn is None or not isinstance(
+                getattr(attn, "q_proj", None), torch.nn.Linear):
+            continue
+        if not _plain_norm(getattr(mod, "norm", None)):
+            continue
+        if not all(_block_ok(b) for b in layers):
+            continue
+        if attn.k_proj.out_features >= attn.q_proj.out_features:
+            continue
+        roots.append(path)
+    return roots
+
+
+def _bind(model, root, probe):
+    return fp8_chain.bind_prefill_fp8_chain(model, root, probe)
+
+
+def _band_candidate(band: str, row: dict) -> regions.RegionCandidate:
+    return regions.RegionCandidate(
+        name=f"{band}_chain",
+        missing=lambda band=band: fp8_chain.missing_symbols(band=band),
+        bind=lambda model, root, probe, band=band:
+            fp8_chain.bind_prefill_fp8_chain(model, root, probe,
+                                             band=band),
+        precision_rank=row["precision_rank"],
+    )
+
+
+#: candidates generate from the band table — a precision band is a
+#: table row in the chain module, never new wiring here
+FAMILY = regions.RegionFamily(
+    family="prefill_tower",
+    identify=identify,
+    candidates=[_band_candidate(band, row)
+                for band, row in fp8_chain.BANDS.items()],
+)
+
+
+def register() -> None:
+    """(Re-)register the family — idempotent, import calls it once."""
+    regions.register_region_family(FAMILY)
+
+
+register()

flashrt_structures/impls/qk_norm_rope/__init__.py

@@ -0,0 +1,12 @@
+from .projection_bf16 import (
+    ProjectionQkNormRope,
+    bind_projection_qk_norm_rope,
+)
+from .per_head_gqa import PerHeadGqaQkNormRope, bind_per_head_gqa_qk_norm_rope
+
+__all__ = [
+    "PerHeadGqaQkNormRope",
+    "ProjectionQkNormRope",
+    "bind_per_head_gqa_qk_norm_rope",
+    "bind_projection_qk_norm_rope",
+]

flashrt_structures/impls/qk_norm_rope/per_head_gqa.py

@@ -0,0 +1,199 @@
+"""Per-head GQA Q/K RMSNorm and rotate-half RoPE implementation."""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardRefused, GuardedSeam
+from ...workspace import lease
+
+
+class PerHeadGqaQkNormRope(GuardedSeam, torch.nn.Module):
+    """Consume packed GQA QKV and produce attention-ready workspaces."""
+
+    _frt_can_fallback = False
+
+    def __init__(
+        self,
+        q_norm_weight: torch.Tensor,
+        k_norm_weight: torch.Tensor,
+        *,
+        row_capacity: int,
+        q_heads: int,
+        kv_heads: int,
+        head_dim: int,
+        eps: float = 1e-6,
+        workspace_lane: str | None = None,
+    ) -> None:
+        super().__init__()
+        if row_capacity <= 0 or q_heads <= 0 or kv_heads <= 0:
+            raise ValueError(
+                "qk_norm_rope: row capacity and head counts must be positive"
+            )
+        if head_dim != 128:
+            raise ValueError(
+                "qk_norm_rope: per-head GQA kernel requires head_dim == 128"
+            )
+        if q_norm_weight.shape != (128,) or k_norm_weight.shape != (128,):
+            raise ValueError(
+                "qk_norm_rope: per-head norm weights must have shape (128,)"
+            )
+        if q_norm_weight.device != k_norm_weight.device:
+            raise ValueError("qk_norm_rope: Q/K norm weights must share device")
+
+        self.row_capacity = int(row_capacity)
+        self.q_heads = int(q_heads)
+        self.kv_heads = int(kv_heads)
+        self.head_dim = 128
+        self.eps = float(eps)
+        kernel = hub_kernel("flashrt/flashrt-qkv-cache-rope", ">=1")
+        try:
+            self._fn = kernel.qkv_split_per_head_norm_rope_bf16
+        except AttributeError as exc:
+            raise ValueError(
+                "qk_norm_rope: flashrt-qkv-cache-rope artifact lacks the "
+                "per-head GQA entry"
+            ) from exc
+        self.register_buffer(
+            "q_norm_weight",
+            q_norm_weight.detach().to(torch.bfloat16).contiguous(),
+        )
+        self.register_buffer(
+            "k_norm_weight",
+            k_norm_weight.detach().to(torch.bfloat16).contiguous(),
+        )
+        device = q_norm_weight.device
+        if workspace_lane is not None:
+            # a caller that declares its outputs call-scoped (no cache,
+            # consumed inside the layer) shares one workspace per lane
+            # across every same-shape layer — the difference between a
+            # 19k-token host binding 52 layers and OOMing on the 53rd
+            self.q_out = lease((self.row_capacity, self.q_heads, 128),
+                               torch.bfloat16, device,
+                               tag=f"qkr_q|{workspace_lane}")
+            self.k_out = lease((self.row_capacity, self.kv_heads, 128),
+                               torch.bfloat16, device,
+                               tag=f"qkr_k|{workspace_lane}")
+            self.v_out = lease((self.row_capacity, self.kv_heads, 128),
+                               torch.bfloat16, device,
+                               tag=f"qkr_v|{workspace_lane}")
+        else:
+            self.register_buffer(
+                "q_out",
+                torch.empty(
+                    self.row_capacity,
+                    self.q_heads,
+                    128,
+                    device=device,
+                    dtype=torch.bfloat16,
+                ),
+                persistent=False,
+            )
+            self.register_buffer(
+                "k_out",
+                torch.empty(
+                    self.row_capacity,
+                    self.kv_heads,
+                    128,
+                    device=device,
+                    dtype=torch.bfloat16,
+                ),
+                persistent=False,
+            )
+            self.register_buffer(
+                "v_out",
+                torch.empty_like(self.k_out),
+                persistent=False,
+            )
+        self._frt_arm(
+            dtypes={torch.bfloat16},
+            device=device,
+            k=(self.q_heads + 2 * self.kv_heads) * 128,
+            row_capacity=self.row_capacity,
+        )
+
+    def forward(
+        self,
+        packed_qkv: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        admitted = self._frt_admit(packed_qkv, cos, sin)
+        if admitted is not PROCEED:
+            return admitted
+        if packed_qkv.dim() != 3:
+            raise GuardRefused(
+                "qk_norm_rope: packed QKV must have shape (B, T, width)"
+            )
+        batch, tokens, width = packed_qkv.shape
+        rows = batch * tokens
+        expected_width = (self.q_heads + 2 * self.kv_heads) * 128
+        if width != expected_width or rows > self.row_capacity:
+            raise GuardRefused(
+                "qk_norm_rope: packed GQA QKV is outside the bound form"
+            )
+        expected_freq = (batch, tokens, 128)
+        if cos.shape != expected_freq or sin.shape != expected_freq:
+            raise GuardRefused(
+                f"qk_norm_rope: cos/sin must have shape {expected_freq}"
+            )
+        if (
+            cos.dtype is not torch.bfloat16
+            or sin.dtype is not torch.bfloat16
+            or cos.device != packed_qkv.device
+            or sin.device != packed_qkv.device
+        ):
+            raise GuardRefused(
+                "qk_norm_rope: cos/sin must be BF16 on the QKV device"
+            )
+        if not packed_qkv.is_contiguous():
+            raise GuardRefused("qk_norm_rope: packed QKV must be contiguous")
+
+        q_out = self.q_out[:rows].view(batch, tokens, self.q_heads, 128)
+        k_out = self.k_out[:rows].view(batch, tokens, self.kv_heads, 128)
+        v_out = self.v_out[:rows].view(batch, tokens, self.kv_heads, 128)
+        return self._fn(
+            packed_qkv,
+            self.q_norm_weight,
+            self.k_norm_weight,
+            cos,
+            sin,
+            self.q_heads,
+            self.kv_heads,
+            eps=self.eps,
+            q_out=q_out,
+            k_out=k_out,
+            v_out=v_out,
+        )
+
+
+def bind_per_head_gqa_qk_norm_rope(
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    *,
+    row_capacity: int,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    eps: float = 1e-6,
+    workspace_lane: str | None = None,
+) -> PerHeadGqaQkNormRope:
+    """Bind the capacity-guarded per-head GQA implementation.
+
+    ``q_norm_weight`` and ``k_norm_weight`` are checkpoint-native vectors
+    with layout ``[head_dim]``.  The packed activation is laid out
+    ``[Q heads..., K heads..., V heads...]``, each head contiguous at
+    ``head_dim`` elements; no transposition or interleaving of the norm
+    weights is performed.
+    """
+    return PerHeadGqaQkNormRope(
+        q_norm_weight,
+        k_norm_weight,
+        row_capacity=row_capacity,
+        q_heads=q_heads,
+        kv_heads=kv_heads,
+        head_dim=head_dim,
+        eps=eps,
+        workspace_lane=workspace_lane,
+    )

flashrt_structures/impls/qk_norm_rope/projection_bf16.py

@@ -0,0 +1,165 @@
+"""Packed-QKV implementation of projection-scope Q/K norm plus RoPE.
+
+This is the Wan form of :mod:`qk_norm_rope`: Q and K are normalized over
+the complete projection before the output is viewed as heads. It consumes
+the contiguous output of a QKV pack and materializes Q/K/V attention
+workspaces in one Hub kernel. Per-head Cosmos/Qwen/audio normalization is a
+different implementation variant of the same catalog structure.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardRefused, GuardedSeam
+
+
+class ProjectionQkNormRope(GuardedSeam, torch.nn.Module):
+    """Fixed-shape packed QKV postprocess for projection-scope RMSNorm."""
+
+    _frt_can_fallback = False
+
+    def __init__(
+        self,
+        q_norm_weight: torch.Tensor,
+        k_norm_weight: torch.Tensor,
+        *,
+        batch: int,
+        tokens: int,
+        heads: int,
+        head_dim: int,
+        qkv_bias: torch.Tensor | None = None,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        if min(batch, tokens, heads, head_dim) <= 0:
+            raise ValueError("qk_norm_rope: shape dimensions must be positive")
+        if head_dim % 2:
+            raise ValueError("qk_norm_rope: head_dim must be even")
+        dim = int(heads) * int(head_dim)
+        if q_norm_weight.numel() != dim or k_norm_weight.numel() != dim:
+            raise ValueError(
+                "qk_norm_rope: projection-scope norm weights must each "
+                f"contain heads * head_dim = {dim} elements")
+        device = q_norm_weight.device
+        if k_norm_weight.device != device:
+            raise ValueError("qk_norm_rope: Q/K norm weights must share device")
+        if qkv_bias is None:
+            qkv_bias = torch.zeros(
+                3 * dim, device=device, dtype=torch.bfloat16)
+        if qkv_bias.numel() != 3 * dim or qkv_bias.device != device:
+            raise ValueError(
+                "qk_norm_rope: qkv_bias must contain 3 * heads * head_dim "
+                "elements on the norm-weight device")
+
+        self.batch = int(batch)
+        self.tokens = int(tokens)
+        self.heads = int(heads)
+        self.head_dim = int(head_dim)
+        self.dim = dim
+        self.eps = float(eps)
+        self._fn = hub_kernel(
+            "flashrt/flashrt-qkv-cache-rope",
+            ">=1",
+        ).qkv_split_bias_norm_rope_v_bf16
+        self.register_buffer(
+            "q_norm_weight",
+            q_norm_weight.detach().reshape(dim).to(torch.bfloat16).contiguous(),
+        )
+        self.register_buffer(
+            "k_norm_weight",
+            k_norm_weight.detach().reshape(dim).to(torch.bfloat16).contiguous(),
+        )
+        self.register_buffer(
+            "qkv_bias",
+            qkv_bias.detach().reshape(3 * dim).to(torch.bfloat16).contiguous(),
+        )
+        shape = (self.batch, self.tokens, self.heads, self.head_dim)
+        self.register_buffer(
+            "q_out",
+            torch.empty(shape, device=device, dtype=torch.bfloat16),
+            persistent=False,
+        )
+        self.register_buffer(
+            "k_out",
+            torch.empty(shape, device=device, dtype=torch.bfloat16),
+            persistent=False,
+        )
+        self.register_buffer(
+            "v_out",
+            torch.empty(shape, device=device, dtype=torch.bfloat16),
+            persistent=False,
+        )
+        self._frt_arm(
+            dtypes={torch.bfloat16},
+            device=device,
+            k=3 * dim,
+            rows=self.batch * self.tokens,
+        )
+
+    def forward(
+        self,
+        packed_qkv: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        admitted = self._frt_admit(packed_qkv, cos, sin)
+        if admitted is not PROCEED:
+            return admitted
+        expected = (self.batch, self.tokens, 3 * self.dim)
+        if tuple(packed_qkv.shape) != expected:
+            raise GuardRefused(
+                f"qk_norm_rope: packed QKV shape {tuple(packed_qkv.shape)} "
+                f"(bound for {expected})")
+        freq_shape = (self.tokens, self.head_dim // 2)
+        if tuple(cos.shape) != freq_shape or tuple(sin.shape) != freq_shape:
+            raise GuardRefused(
+                f"qk_norm_rope: cos/sin must have shape {freq_shape}")
+        if (cos.dtype is not torch.float32 or sin.dtype is not torch.float32
+                or cos.device != packed_qkv.device
+                or sin.device != packed_qkv.device):
+            raise GuardRefused(
+                "qk_norm_rope: cos/sin must be float32 on the QKV device")
+        if not packed_qkv.is_contiguous():
+            raise GuardRefused("qk_norm_rope: packed QKV must be contiguous")
+
+        return self._fn(
+            packed_qkv,
+            self.qkv_bias,
+            self.q_norm_weight,
+            self.k_norm_weight,
+            cos,
+            sin,
+            self.heads,
+            self.head_dim,
+            rope_seq_len=self.tokens,
+            eps=self.eps,
+            q_out=self.q_out,
+            k_out=self.k_out,
+            v_out=self.v_out,
+        )
+
+
+def bind_projection_qk_norm_rope(
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    *,
+    batch: int,
+    tokens: int,
+    heads: int,
+    head_dim: int,
+    qkv_bias: torch.Tensor | None = None,
+    eps: float = 1e-6,
+) -> ProjectionQkNormRope:
+    """Bind the fixed-shape projection-scope implementation."""
+    return ProjectionQkNormRope(
+        q_norm_weight,
+        k_norm_weight,
+        batch=batch,
+        tokens=tokens,
+        heads=heads,
+        head_dim=head_dim,
+        qkv_bias=qkv_bias,
+        eps=eps,
+    )

flashrt_structures/impls/qkv_pack/__init__.py

@@ -0,0 +1,5 @@
+from .fp8_static import (AttnBlockPacked, PackedLinear, StashReader,
+                         bind_attn_block, bind_qkv_pack)
+
+__all__ = ["AttnBlockPacked", "PackedLinear", "StashReader",
+           "bind_attn_block", "bind_qkv_pack"]

flashrt_structures/impls/qkv_pack/bf16.py

@@ -0,0 +1,110 @@
+"""Non-quantising BF16 implementation of ``qkv_pack``.
+
+Sibling projections that proved shared-input fixed-order dataflow are one
+larger BF16 GEMM.  No quantisation is introduced: this is the portable
+structural form used when an end-to-end accuracy gate refuses FP8/FP4.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from ...guard import CAST_OK, PROCEED, GuardedSeam
+
+
+class PackedBf16Linear(GuardedSeam, torch.nn.Module):
+    _frt_host_attr = "host_linear"
+    _frt_can_fallback = True
+
+    def __init__(self, mods: Sequence[torch.nn.Linear], rows: int):
+        super().__init__()
+        if len(mods) < 2:
+            raise ValueError("qkv_pack: need at least two siblings")
+        kdims = {int(mod.weight.shape[1]) for mod in mods}
+        if len(kdims) != 1:
+            raise ValueError(f"qkv_pack: sibling K dims differ {kdims}")
+        self.splits = tuple(int(mod.weight.shape[0]) for mod in mods)
+        weight = torch.cat(
+            [mod.weight.detach() for mod in mods], dim=0).contiguous()
+        bias = None
+        if any(mod.bias is not None for mod in mods):
+            bias = torch.cat([
+                (mod.bias.detach() if mod.bias is not None else
+                 torch.zeros(mod.weight.shape[0], device=mod.weight.device,
+                             dtype=mod.weight.dtype))
+                for mod in mods
+            ]).contiguous()
+        self.register_buffer("packed_weight", weight)
+        self.register_buffer("packed_bias", bias)
+        self.host_linear = mods[0]
+        for index, width in enumerate(self.splits[1:], 1):
+            self.register_buffer(
+                f"stash{index}", torch.empty(
+                    rows, width, device=weight.device, dtype=weight.dtype))
+        self._frt_arm(dtypes=CAST_OK, device=weight.device,
+                      k=next(iter(kdims)), row_capacity=rows)
+
+    def _run(self, flat: torch.Tensor) -> torch.Tensor:
+        y = torch.nn.functional.linear(
+            flat.to(self.packed_weight.dtype), self.packed_weight,
+            self.packed_bias)
+        offset = self.splits[0]
+        for index, width in enumerate(self.splits[1:], 1):
+            getattr(self, f"stash{index}")[:flat.shape[0]].copy_(
+                y[:, offset:offset + width])
+            offset += width
+        return y
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        y = self._run(x.reshape(-1, x.shape[-1]))
+        out = y[:, :self.splits[0]].contiguous()
+        return out.reshape(*x.shape[:-1], self.splits[0]).to(x.dtype)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(super().__getattr__("host_linear"), name)
+
+
+class Bf16StashReader(GuardedSeam, torch.nn.Module):
+    _frt_host_attr = "host_linear"
+    _frt_can_fallback = True
+    _frt_requires_sibling_order = True
+
+    def __init__(self, original: torch.nn.Linear,
+                 packed: PackedBf16Linear, index: int):
+        super().__init__()
+        self.host_linear = original
+        self._packed = (packed,)
+        self.index = int(index)
+        head = packed._frt_guard
+        self._frt_arm(dtypes=head.dtypes, device=head.device, k=head.k,
+                      row_capacity=head.row_capacity)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        admitted = self._frt_admit(x)
+        if admitted is not PROCEED:
+            return admitted
+        rows = x.numel() // x.shape[-1]
+        out = getattr(self._packed[0], f"stash{self.index}")[:rows]
+        return out.reshape(*x.shape[:-1], out.shape[-1]).to(x.dtype)
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(super().__getattr__("host_linear"), name)
+
+
+def bind_qkv_pack(mods: Sequence[torch.nn.Linear], *, rows: int):
+    packed = PackedBf16Linear(mods, rows)
+    return [packed, *(
+        Bf16StashReader(mod, packed, index)
+        for index, mod in enumerate(mods[1:], 1)
+    )]