stackformers 3.8.0__py3-none-any.whl

stackformers/attention/self_attn.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch import Tensor
+
+from stackformers.attention.bias import NoAttnBias
+from stackformers.attention.config import SelfAttentionConfig
+from stackformers.attention.ops import packed_attn_or_fallback, padded_sdpa
+from stackformers.attention.protocols import AttnBias
+from stackformers.positional.protocols import PosEncoding
+from stackformers.sequence import (
+    PackedInput,
+    PackedSequence,
+    PaddedInput,
+    SequenceInput,
+    packed_to_padded,
+)
+
+
+class SelfAttention(nn.Module):
+    """Multi-head self-attention.
+
+    window_size=None (default) → global attention.
+    window_size=w             → sliding-window attention with width w.
+
+    Pass PaddedInput for inference or export, PackedInput for training — same weights.
+    Falls back to padded SDPA when varlen_attn is unavailable (CPU, float32, torch.export)
+    or when an attention bias is provided.
+    """
+
+    def __init__(
+        self,
+        config: SelfAttentionConfig,
+        pos_encoding: PosEncoding,
+        attn_bias: AttnBias = NoAttnBias(),
+    ) -> None:
+        super().__init__()
+        self.config = config
+        h, kv_h, dh = config.heads, config.effective_kv_heads, config.dim_head
+        self.to_q = nn.Linear(config.dim, h * dh, bias=False)
+        self.to_k = nn.Linear(config.dim, kv_h * dh, bias=False)
+        self.to_v = nn.Linear(config.dim, kv_h * dh, bias=False)
+        self.to_out = nn.Linear(h * dh, config.dim, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.pos_encoding = pos_encoding
+        self.attn_bias = attn_bias
+        self.q_norm: nn.Module = nn.RMSNorm(dh) if config.qk_norm else nn.Identity()
+        self.k_norm: nn.Module = nn.RMSNorm(dh) if config.qk_norm else nn.Identity()
+        nn.init.normal_(self.to_out.weight, std=0.02)
+
+    def _forward_padded(self, input: PaddedInput) -> Tensor:
+        cfg = self.config
+        h, kv_h, groups = cfg.heads, cfg.effective_kv_heads, cfg.groups
+        x = input.x
+        q = self.q_norm(rearrange(self.to_q(x), "b n (h d) -> b h n d", h=h))
+        k = self.k_norm(rearrange(self.to_k(x), "b n (h d) -> b h n d", h=kv_h))
+        v = rearrange(self.to_v(x), "b n (h d) -> b h n d", h=kv_h)
+        if groups > 1:
+            k = repeat(k, "b h n d -> b (h g) n d", g=groups)
+            v = repeat(v, "b h n d -> b (h g) n d", g=groups)
+        q, k = self.pos_encoding.forward_padded(q, k, input.abs_positions, input.abs_positions)
+        bias = self.attn_bias(input)
+        out = padded_sdpa(q, k, v, input.mask, cfg.causal, cfg.window_size, bias)
+        return self.dropout(self.to_out(rearrange(out, "b h n d -> b n (h d)")))
+
+    def _forward_packed(self, input: PackedInput) -> Tensor:
+        cfg = self.config
+        h, kv_h, groups = cfg.heads, cfg.effective_kv_heads, cfg.groups
+        x = input.x
+        q = self.q_norm(rearrange(self.to_q(x), "nt (h d) -> nt h d", h=h))
+        k = self.k_norm(rearrange(self.to_k(x), "nt (h d) -> nt h d", h=kv_h))
+        v = rearrange(self.to_v(x), "nt (h d) -> nt h d", h=kv_h)
+        if groups > 1:
+            k = repeat(k, "nt h d -> nt (h g) d", g=groups)
+            v = repeat(v, "nt h d -> nt (h g) d", g=groups)
+        q, k = self.pos_encoding.forward_packed(q, k, input.abs_positions, input.abs_positions)
+        bias = self.attn_bias(packed_to_padded(input))
+        seq = PackedSequence(cu_seqlens=input.cu_seqlens, max_seqlen=input.max_seqlen)
+        out = packed_attn_or_fallback(q, k, v, seq, seq, cfg.causal, cfg.window_size, bias)
+        return self.dropout(self.to_out(rearrange(out, "nt h d -> nt (h d)")))
+
+    def forward(self, input: SequenceInput) -> Tensor:
+        match input:
+            case PaddedInput():
+                return self._forward_padded(input)
+            case PackedInput():
+                return self._forward_packed(input)

stackformers/config.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from stackformers.attention.config import CrossAttentionConfig, SelfAttentionConfig
+from stackformers.feedforward.config import FeedForwardConfig
+
+
+class LayerConfig(BaseModel):
+    attn: SelfAttentionConfig
+    ff: FeedForwardConfig
+    pre_norm: bool = True
+
+
+class EncoderConfig(BaseModel):
+    layer: LayerConfig
+    num_layers: int = Field(gt=0)
+    dropout: float = Field(default=0.0, ge=0.0, le=1.0)
+
+
+class DecoderConfig(BaseModel):
+    self_attn: SelfAttentionConfig
+    cross_attn: CrossAttentionConfig
+    ff: FeedForwardConfig
+    num_layers: int = Field(gt=0)
+    dropout: float = Field(default=0.0, ge=0.0, le=1.0)

stackformers/cross_attender.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from torch import Tensor
+
+from stackformers.attention.protocols import CrossAttn
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.protocols import Norm
+from stackformers.sequence import SequenceInput
+
+
+class CrossAttenderLayer(nn.Module):
+    """Pre-norm layer: cross-attn → feed-forward. No self-attention."""
+
+    def __init__(
+        self,
+        cross_attn: CrossAttn,
+        ff: FeedForward,
+        norm_cross: Norm,
+        norm_ff: Norm,
+    ) -> None:
+        super().__init__()
+        self.cross_attn = cross_attn
+        self.ff = ff
+        self.norm_cross = norm_cross
+        self.norm_ff = norm_ff
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> SequenceInput:
+        normed = x_input._replace(x=self.norm_cross(x_input.x))
+        x = x_input.x + self.cross_attn(normed, ctx_input)
+        x = x + self.ff(self.norm_ff(x))
+        return x_input._replace(x=x)
+
+
+class CrossAttenderStack(nn.Module):
+    """Stack of CrossAttenderLayers with a final layer norm."""
+
+    def __init__(self, layers: list[CrossAttenderLayer], final_norm: Norm) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+        self.final_norm = final_norm
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor:
+        for layer in self.layers:
+            x_input = layer(x_input, ctx_input)
+        return self.final_norm(x_input.x)

stackformers/decoder.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from torch import Tensor
+
+from stackformers.attention.protocols import CrossAttn, SelfAttn
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.protocols import Norm
+from stackformers.sequence import SequenceInput
+
+
+class DecoderLayer(nn.Module):
+    """Pre-norm decoder layer: causal self-attn → cross-attn → feed-forward."""
+
+    def __init__(
+        self,
+        self_attn: SelfAttn,
+        cross_attn: CrossAttn,
+        ff: FeedForward,
+        norm_self: Norm,
+        norm_cross: Norm,
+        norm_ff: Norm,
+    ) -> None:
+        super().__init__()
+        self.self_attn = self_attn
+        self.cross_attn = cross_attn
+        self.ff = ff
+        self.norm_self = norm_self
+        self.norm_cross = norm_cross
+        self.norm_ff = norm_ff
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> SequenceInput:
+        normed_self = x_input._replace(x=self.norm_self(x_input.x))
+        x = x_input.x + self.self_attn(normed_self)
+        x_input = x_input._replace(x=x)
+        normed_cross = x_input._replace(x=self.norm_cross(x_input.x))
+        x = x_input.x + self.cross_attn(normed_cross, ctx_input)
+        x = x + self.ff(self.norm_ff(x))
+        return x_input._replace(x=x)
+
+
+class Decoder(nn.Module):
+    """Stack of DecoderLayers with a final layer norm."""
+
+    def __init__(
+        self,
+        layers: list[DecoderLayer],
+        final_norm: Norm,
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+        self.final_norm = final_norm
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor:
+        for layer in self.layers:
+            x_input = layer(x_input, ctx_input)
+        return self.final_norm(x_input.x)

stackformers/encoder.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from torch import Tensor
+
+from stackformers.layers import TransformerLayer
+from stackformers.norm.protocols import Norm
+from stackformers.sequence import SequenceInput
+
+
+class Encoder(nn.Module):
+    """Stack of TransformerLayers with a final layer norm."""
+
+    def __init__(
+        self,
+        layers: list[TransformerLayer],
+        final_norm: Norm,
+    ) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(layers)
+        self.final_norm = final_norm
+
+    def forward(self, input: SequenceInput) -> Tensor:
+        for layer in self.layers:
+            input = layer(input)
+        return self.final_norm(input.x)

stackformers/feedforward/README.md

@@ -0,0 +1,11 @@
+# feedforward
+
+Token-wise feed-forward sublayers behind the `FeedForward` protocol: `(x: b n d) → b n d`.
+
+`SwiGLU` is the only implementation. Its inner dimension is `int(dim * mult * 2/3)`, scaled down so parameter count matches a standard 4× GELU FFN.
+
+## Adding a new feed-forward
+
+1. Add a config class to `config.py` (or reuse `FeedForwardConfig` if parameterisation is identical).
+2. Implement the class satisfying `FeedForward` structurally.
+3. Add a `case` branch in `factory.py::build_ff`.

stackformers/feedforward/__init__.py

@@ -0,0 +1,5 @@
+from stackformers.feedforward.geglu import GEGLU
+from stackformers.feedforward.relu_squared import ReluSquaredFF
+from stackformers.feedforward.swiglu import SwiGLU
+
+__all__ = ["GEGLU", "ReluSquaredFF", "SwiGLU"]

stackformers/feedforward/config.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import warnings
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class _FFBase(BaseModel):
+    """Shared geometry and alignment check for gated feed-forward variants."""
+
+    dim: int = Field(gt=0)
+    mult: float = Field(default=4.0, gt=0.0)
+    dropout: float = Field(default=0.0, ge=0.0, le=1.0)
+
+    _ALIGN = 64  # tensor-core alignment for FP16/BF16
+
+    @model_validator(mode="after")
+    def _check_inner_dim_alignment(self) -> _FFBase:
+        d = self.inner_dim
+        if d % self._ALIGN != 0:
+            warnings.warn(
+                f"inner_dim={d} is not a multiple of {self._ALIGN}. "
+                "Unaligned dimensions reduce GPU throughput on tensor-core hardware. "
+                f"Nearest aligned values: {(d // self._ALIGN) * self._ALIGN} or "
+                f"{(d // self._ALIGN + 1) * self._ALIGN}. "
+                "Adjust dim or mult so that int(dim * mult * 2/3) is a multiple of "
+                f"{self._ALIGN}.",
+                UserWarning,
+                stacklevel=2,
+            )
+        return self
+
+    @property
+    def inner_dim(self) -> int:
+        # Two gate matrices → scale inner dim down to match param count with GELU-4x FFN.
+        return int(self.dim * self.mult * 2 / 3)
+
+
+class SwiGLUConfig(_FFBase):
+    """Config for the SwiGLU feed-forward network (Noam Shazeer, 2020)."""
+
+    kind: Literal["swiglu"] = "swiglu"
+
+
+class GEGLUConfig(_FFBase):
+    """Config for the GEGLU feed-forward network (Noam Shazeer, 2020)."""
+
+    kind: Literal["geglu"] = "geglu"
+
+
+class ReluSquaredConfig(_FFBase):
+    """Config for the ReLU² feed-forward network.
+
+    Non-gated: inner_dim = dim * mult (no 2/3 factor).
+    Paper: "Primer: Searching for Efficient Transformers" — https://arxiv.org/abs/2109.08668
+    """
+
+    kind: Literal["relu_squared"] = "relu_squared"
+
+    @property
+    def inner_dim(self) -> int:
+        return int(self.dim * self.mult)
+
+
+FeedForwardConfig = Annotated[
+    SwiGLUConfig | GEGLUConfig | ReluSquaredConfig, Field(discriminator="kind")
+]

stackformers/feedforward/factory.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from stackformers.feedforward.config import (
+    FeedForwardConfig,
+    GEGLUConfig,
+    ReluSquaredConfig,
+    SwiGLUConfig,
+)
+from stackformers.feedforward.geglu import GEGLU
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.feedforward.relu_squared import ReluSquaredFF
+from stackformers.feedforward.swiglu import SwiGLU
+
+
+def build_ff(config: FeedForwardConfig) -> FeedForward:
+    match config:
+        case SwiGLUConfig():
+            return SwiGLU(config)
+        case GEGLUConfig():
+            return GEGLU(config)
+        case ReluSquaredConfig():
+            return ReluSquaredFF(config)
+        case _:
+            raise AssertionError(f"Unhandled feedforward config: {type(config)}")

stackformers/feedforward/geglu.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.feedforward.config import GEGLUConfig
+
+
+class GEGLU(nn.Module):
+    """GEGLU feed-forward network.
+
+    GLU variant that gates with GELU instead of SiLU (SwiGLU). Empirically
+    competitive with SwiGLU; slightly smoother gradient signal near zero.
+    inner_dim = 2/3 * dim * mult so total params match a standard 4x GELU FFN.
+
+    Paper: "GLU Variants Improve Transformers" — https://arxiv.org/abs/2002.05202
+    """
+
+    def __init__(self, config: GEGLUConfig) -> None:
+        super().__init__()
+        inner_dim = config.inner_dim
+
+        self.w1 = nn.Linear(config.dim, inner_dim, bias=False)
+        self.w2 = nn.Linear(config.dim, inner_dim, bias=False)
+        self.w3 = nn.Linear(inner_dim, config.dim, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.act = nn.GELU(approximate="tanh")
+
+        nn.init.normal_(self.w3.weight, std=0.02)
+
+    def forward(self, x: Float[Tensor, "b n d"]) -> Float[Tensor, "b n d"]:
+        gate = self.act(self.w1(x))
+        hidden = gate * self.w2(x)
+        return self.w3(self.dropout(hidden))

stackformers/feedforward/protocols.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Float
+from torch import Tensor
+
+
+@runtime_checkable
+class FeedForward(Protocol):
+    """Apply a feed-forward transformation to (b, n, d) token embeddings.
+
+    Implementation: SwiGLU.
+    """
+
+    def __call__(self, x: Float[Tensor, "b n d"]) -> Float[Tensor, "b n d"]: ...

stackformers/feedforward/relu_squared.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.feedforward.config import ReluSquaredConfig
+
+
+class ReluSquaredFF(nn.Module):
+    """ReLU² feed-forward network.
+
+    Standard two-matrix FFN (no gating) with ReLU² activation.
+    inner_dim = dim * mult; parameter count matches a standard mult-x FFN.
+
+    Paper: "Primer: Searching for Efficient Transformers" — https://arxiv.org/abs/2109.08668
+    """
+
+    def __init__(self, config: ReluSquaredConfig) -> None:
+        super().__init__()
+        self.w1 = nn.Linear(config.dim, config.inner_dim, bias=False)
+        self.w2 = nn.Linear(config.inner_dim, config.dim, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.act = nn.ReLU()
+
+    def forward(self, x: Float[Tensor, "b n d"]) -> Float[Tensor, "b n d"]:
+        h = self.act(self.w1(x))
+        return self.w2(self.dropout(h.square()))

stackformers/feedforward/swiglu.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.feedforward.config import SwiGLUConfig
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU feed-forward network.
+
+    Uses two parallel gate matrices (W1, W2) and one output projection (W3).
+    inner_dim = 2/3 * dim * mult so total params match a standard 4x GELU FFN.
+
+    Paper: "GLU Variants Improve Transformers" — https://arxiv.org/abs/2002.05202
+    """
+
+    def __init__(self, config: SwiGLUConfig) -> None:
+        super().__init__()
+        inner_dim = config.inner_dim
+
+        self.w1 = nn.Linear(config.dim, inner_dim, bias=False)
+        self.w2 = nn.Linear(config.dim, inner_dim, bias=False)
+        self.w3 = nn.Linear(inner_dim, config.dim, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.act = nn.SiLU()
+
+        nn.init.normal_(self.w3.weight, std=0.02)
+
+    def forward(self, x: Float[Tensor, "b n d"]) -> Float[Tensor, "b n d"]:
+        gate = self.act(self.w1(x))
+        hidden = gate * self.w2(x)
+        return self.w3(self.dropout(hidden))

stackformers/layers.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import torch.nn as nn
+
+from stackformers.attention.protocols import SelfAttn
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.protocols import Norm
+from stackformers.sequence import SequenceInput
+
+
+class TransformerLayer(nn.Module):
+    """Pre-norm transformer layer: norm → self-attn → residual, norm → ff → residual."""
+
+    def __init__(
+        self,
+        self_attn: SelfAttn,
+        ff: FeedForward,
+        norm_attn: Norm,
+        norm_ff: Norm,
+    ) -> None:
+        super().__init__()
+        self.self_attn = self_attn
+        self.ff = ff
+        self.norm_attn = norm_attn
+        self.norm_ff = norm_ff
+
+    def forward(self, input: SequenceInput) -> SequenceInput:
+        normed = input._replace(x=self.norm_attn(input.x))
+        x = input.x + self.self_attn(normed)
+        x = x + self.ff(self.norm_ff(x))
+        return input._replace(x=x)

stackformers/norm/README.md

@@ -0,0 +1,11 @@
+# norm
+
+Normalisation layers behind the `Norm` protocol: `(x: b n d) → b n d`.
+
+`RMSNorm` is the default — no mean subtraction, no bias, numerically stable on fp16. `LayerNorm` is available for compatibility.
+
+## Adding a new norm
+
+1. Add a config class with a `kind: Literal[...]` discriminator to `config.py` and include it in the `NormConfig` union.
+2. Implement the class satisfying `Norm` structurally — do not import the protocol.
+3. Add a `case` branch in `factory.py::build_norm`.

stackformers/norm/__init__.py

@@ -0,0 +1,3 @@
+from stackformers.norm.config import LayerNormConfig, RMSNormConfig
+
+__all__ = ["LayerNormConfig", "RMSNormConfig"]

stackformers/norm/config.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+
+class RMSNormConfig(BaseModel):
+    kind: Literal["rms"] = "rms"
+    dim: int = Field(gt=0)
+    eps: float = Field(default=1e-6, gt=0)
+
+
+class LayerNormConfig(BaseModel):
+    kind: Literal["layer_norm"] = "layer_norm"
+    dim: int = Field(gt=0)
+    eps: float = 1e-5
+
+
+NormConfig = Annotated[RMSNormConfig | LayerNormConfig, Field(discriminator="kind")]

stackformers/norm/factory.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import torch.nn as nn
+
+from stackformers.norm.config import LayerNormConfig, NormConfig, RMSNormConfig
+from stackformers.norm.protocols import Norm
+
+__all__ = ["NormConfig", "build_norm"]
+
+
+def build_norm(config: NormConfig) -> Norm:
+    match config:
+        case RMSNormConfig():
+            return nn.RMSNorm(config.dim, eps=config.eps)
+        case LayerNormConfig():
+            return nn.LayerNorm(config.dim, eps=config.eps)
+        case _:
+            raise AssertionError(f"Unhandled norm config: {type(config)}")

stackformers/norm/protocols.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Float
+from torch import Tensor
+
+
+@runtime_checkable
+class Norm(Protocol):
+    """Apply layer normalisation to (b, n, d) tensors.
+
+    Implementation: RMSNorm.
+    """
+
+    def __call__(self, x: Float[Tensor, "b n d"]) -> Float[Tensor, "b n d"]: ...

stackformers/positional/README.md

@@ -0,0 +1,13 @@
+# positional
+
+Positional encodings applied to Q and K tensors inside attention — not to the residual stream.
+
+`RotaryEmbedding1D` is the standard choice. It supports YaRN context extension via `RoPE1DConfig(yarn=YaRNConfig(...))`. `RotaryEmbedding2D` handles grid inputs by splitting the head dimension into row and column components. `NoPosEncoding` is a null object for cross-attention paths that need no positional information.
+
+The `PosEncoding` protocol exposes two methods — `forward_padded` and `forward_packed` — so callers (which already know their layout) pick the right path directly without internal dispatch. All implementations share the same computation via layout-agnostic helpers, since positional encoding is a per-token operation that does not depend on sequence boundaries.
+
+## Adding a new encoding
+
+1. Add a config class with a `kind: Literal[...]` discriminator to `config.py` and include it in the `PosEncodingConfig` union.
+2. Implement the class satisfying `PosEncoding` (and optionally `PackedPosEncoding`) structurally.
+3. Add a `case` branch in `factory.py::build_pos_encoding`.

stackformers/positional/__init__.py

@@ -0,0 +1,23 @@
+from stackformers.positional.config import (
+    LearnedPosEncodingConfig,
+    NoPosEncodingConfig,
+    PosEncodingConfig,
+    RoPE1DConfig,
+    YaRNConfig,
+)
+from stackformers.positional.learned import LearnedPosEncoding
+from stackformers.positional.none import NoPosEncoding
+from stackformers.positional.rope1d import RotaryEmbedding1D
+from stackformers.positional.rope2d import RotaryEmbedding2D
+
+__all__ = [
+    "LearnedPosEncoding",
+    "LearnedPosEncodingConfig",
+    "NoPosEncoding",
+    "NoPosEncodingConfig",
+    "PosEncodingConfig",
+    "RoPE1DConfig",
+    "RotaryEmbedding1D",
+    "RotaryEmbedding2D",
+    "YaRNConfig",
+]