stackformers 3.8.0__py3-none-any.whl

stackformers/presets/decoder.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+import torch.nn as nn
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from stackformers.attention.bias import NoAttnBias
+from stackformers.attention.config import CrossAttentionConfig, SelfAttentionConfig
+from stackformers.attention.cross_attn import CrossAttention
+from stackformers.attention.protocols import AttnBias
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.decoder import Decoder, DecoderLayer
+from stackformers.feedforward.config import FeedForwardConfig, SwiGLUConfig
+from stackformers.feedforward.factory import build_ff
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import NormConfig, build_norm
+from stackformers.norm.protocols import Norm
+from stackformers.positional.config import NoPosEncodingConfig, PosEncodingConfig, RoPE1DConfig
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.positional.none import NoPosEncoding
+from stackformers.positional.protocols import PosEncoding
+from stackformers.sequence import SequenceInput
+
+C = TypeVar("C")
+
+
+class TransformerDecoderBase(nn.Module, Generic[C], ABC):
+    """Abstract decoder: causal self-attn → cross-attn → ff per layer.
+
+    Subclass with any config type C.  Implement build_layers and build_norm;
+    __init__ wires them into a Decoder and nothing else.
+    """
+
+    def __init__(self, config: C) -> None:
+        super().__init__()
+        self.config = config
+        self._decoder = Decoder(
+            layers=self.build_layers(config),
+            final_norm=self.build_norm(config),
+        )
+
+    @abstractmethod
+    def build_layers(self, config: C) -> list[DecoderLayer]: ...
+
+    @abstractmethod
+    def build_norm(self, config: C) -> Norm: ...
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor:
+        return self._decoder(x_input, ctx_input)
+
+
+class TransformerDecoderConfig(BaseModel):
+    self_attn: SelfAttentionConfig
+    cross_attn: CrossAttentionConfig
+    ff: FeedForwardConfig
+    norm: NormConfig
+    pos_encoding: PosEncodingConfig  # applies to self-attention only
+    num_layers: int = Field(gt=0)
+
+
+def plain_decoder_config(
+    dim: int,
+    heads: int,
+    num_layers: int,
+    *,
+    ff_mult: float = 4.0,
+    dropout: float = 0.0,
+) -> TransformerDecoderConfig:
+    """Causal self-attn + cross-attn decoder with RoPE-1D, RMSNorm, SwiGLU FF."""
+    dim_head = dim // heads
+    return TransformerDecoderConfig(
+        self_attn=SelfAttentionConfig(
+            dim=dim, heads=heads, dim_head=dim_head, causal=True, dropout=dropout
+        ),
+        cross_attn=CrossAttentionConfig(dim=dim, heads=heads, dim_head=dim_head, dropout=dropout),
+        ff=SwiGLUConfig(dim=dim, mult=ff_mult, dropout=dropout),
+        norm=RMSNormConfig(dim=dim),
+        pos_encoding=RoPE1DConfig(dim_head=dim_head),
+        num_layers=num_layers,
+    )
+
+
+class TransformerDecoder(TransformerDecoderBase[TransformerDecoderConfig]):
+    """Concrete decoder for TransformerDecoderConfig.
+
+    Override build_self_pos_encoding, build_self_attn_bias, build_ff, or build_norm
+    to customise individual collaborators while keeping the rest of the defaults.
+    """
+
+    def build_layers(self, config: TransformerDecoderConfig) -> list[DecoderLayer]:
+        self_pos = self.build_self_pos_encoding(config)
+        self_bias = self.build_self_attn_bias(config)
+        return [
+            DecoderLayer(
+                self_attn=SelfAttention(
+                    config=config.self_attn,
+                    pos_encoding=self_pos,
+                    attn_bias=self_bias,
+                ),
+                cross_attn=CrossAttention(
+                    config=config.cross_attn,
+                    pos_encoding=NoPosEncoding(NoPosEncodingConfig()),
+                ),
+                ff=self.build_ff(config),
+                norm_self=build_norm(config.norm),
+                norm_cross=build_norm(config.norm),
+                norm_ff=build_norm(config.norm),
+            )
+            for _ in range(config.num_layers)
+        ]
+
+    def build_self_pos_encoding(self, config: TransformerDecoderConfig) -> PosEncoding:
+        return build_pos_encoding(config.pos_encoding)
+
+    def build_self_attn_bias(self, _: TransformerDecoderConfig) -> AttnBias:
+        return NoAttnBias()
+
+    def build_ff(self, config: TransformerDecoderConfig) -> FeedForward:
+        return build_ff(config.ff)
+
+    def build_norm(self, config: TransformerDecoderConfig) -> Norm:
+        return build_norm(config.norm)

stackformers/presets/encoder.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+import torch.nn as nn
+from pydantic import BaseModel, Field
+from torch import Tensor
+
+from stackformers.attention.bias import NoAttnBias
+from stackformers.attention.config import SelfAttentionConfig
+from stackformers.attention.protocols import AttnBias
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.encoder import Encoder
+from stackformers.feedforward.config import FeedForwardConfig, SwiGLUConfig
+from stackformers.feedforward.factory import build_ff
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.layers import TransformerLayer
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import NormConfig, build_norm
+from stackformers.norm.protocols import Norm
+from stackformers.positional.config import PosEncodingConfig, RoPE1DConfig
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.positional.protocols import PosEncoding
+from stackformers.sequence import SequenceInput
+
+C = TypeVar("C")
+
+
+class TransformerEncoderBase(nn.Module, Generic[C], ABC):
+    """Abstract encoder: self-attn → ff per layer.
+
+    Subclass with any config type C.  Implement build_layers and build_norm;
+    __init__ wires them into an Encoder and nothing else.
+    """
+
+    def __init__(self, config: C) -> None:
+        super().__init__()
+        self.config = config
+        self._encoder = Encoder(
+            layers=self.build_layers(config),
+            final_norm=self.build_norm(config),
+        )
+
+    @abstractmethod
+    def build_layers(self, config: C) -> list[TransformerLayer]: ...
+
+    @abstractmethod
+    def build_norm(self, config: C) -> Norm: ...
+
+    def forward(self, input: SequenceInput) -> Tensor:
+        return self._encoder(input)
+
+
+class TransformerEncoderConfig(BaseModel):
+    attn: SelfAttentionConfig
+    ff: FeedForwardConfig
+    norm: NormConfig
+    pos_encoding: PosEncodingConfig
+    num_layers: int = Field(gt=0)
+
+
+def plain_encoder_config(
+    dim: int,
+    heads: int,
+    num_layers: int,
+    *,
+    causal: bool = False,
+    ff_mult: float = 4.0,
+    dropout: float = 0.0,
+) -> TransformerEncoderConfig:
+    """Global SDPA encoder with RoPE-1D, RMSNorm, and SwiGLU FF.
+
+    Pass PaddedInput for inference, PackedInput for training — same model.
+    """
+    dim_head = dim // heads
+    return TransformerEncoderConfig(
+        attn=SelfAttentionConfig(
+            dim=dim, heads=heads, dim_head=dim_head, causal=causal, dropout=dropout
+        ),
+        ff=SwiGLUConfig(dim=dim, mult=ff_mult, dropout=dropout),
+        norm=RMSNormConfig(dim=dim),
+        pos_encoding=RoPE1DConfig(dim_head=dim_head),
+        num_layers=num_layers,
+    )
+
+
+def windowed_encoder_config(
+    dim: int,
+    heads: int,
+    num_layers: int,
+    window_size: int,
+    *,
+    causal: bool = False,
+    ff_mult: float = 4.0,
+    dropout: float = 0.0,
+) -> TransformerEncoderConfig:
+    """Sliding-window encoder — O(n·w) attention for long sequences.
+
+    Pass PaddedInput for inference, PackedInput for training — same model.
+    """
+    dim_head = dim // heads
+    return TransformerEncoderConfig(
+        attn=SelfAttentionConfig(
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            causal=causal,
+            dropout=dropout,
+            window_size=window_size,
+        ),
+        ff=SwiGLUConfig(dim=dim, mult=ff_mult, dropout=dropout),
+        norm=RMSNormConfig(dim=dim),
+        pos_encoding=RoPE1DConfig(dim_head=dim_head),
+        num_layers=num_layers,
+    )
+
+
+class TransformerEncoder(TransformerEncoderBase[TransformerEncoderConfig]):
+    """Concrete encoder for TransformerEncoderConfig.
+
+    Pass PaddedInput for inference, PackedInput for training.
+    Override build_pos_encoding, build_attn_bias, build_ff, or build_norm to customise
+    individual collaborators while keeping the rest of the defaults.
+    """
+
+    def build_layers(self, config: TransformerEncoderConfig) -> list[TransformerLayer]:
+        pos = self.build_pos_encoding(config)
+        bias = self.build_attn_bias(config)
+        return [
+            TransformerLayer(
+                self_attn=SelfAttention(config=config.attn, pos_encoding=pos, attn_bias=bias),
+                ff=self.build_ff(config),
+                norm_attn=build_norm(config.norm),
+                norm_ff=build_norm(config.norm),
+            )
+            for _ in range(config.num_layers)
+        ]
+
+    def build_pos_encoding(self, config: TransformerEncoderConfig) -> PosEncoding:
+        return build_pos_encoding(config.pos_encoding)
+
+    def build_attn_bias(self, _: TransformerEncoderConfig) -> AttnBias:
+        return NoAttnBias()
+
+    def build_ff(self, config: TransformerEncoderConfig) -> FeedForward:
+        return build_ff(config.ff)
+
+    def build_norm(self, config: TransformerEncoderConfig) -> Norm:
+        return build_norm(config.norm)

stackformers/sequence.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+from jaxtyping import Bool, Float, Int
+from torch import Tensor
+
+# NamedTuple instead of dataclass: PyTorch's pytree system handles tuples natively,
+# so torch.export and torch.compile see the tensor fields without any registration.
+# Dataclasses require explicit pytree.register_dataclass() to achieve the same.
+
+
+class PaddedSequence(NamedTuple):
+    """Batch of padded sequences; mask marks valid (non-padding) positions."""
+
+    mask: Bool[Tensor, "b n"]  # True = valid token
+
+
+class PackedSequence(NamedTuple):
+    """Variable-length sequences packed into a single flat tensor.
+
+    cu_seqlens[i] is the cumulative token count up to (not including) sequence i.
+    Follows the FlashAttention convention: length[i] = cu_seqlens[i+1] - cu_seqlens[i].
+    """
+
+    cu_seqlens: Int[Tensor, "bp1"]  # shape (batch + 1,)
+    max_seqlen: int
+
+
+SequenceInfo = PaddedSequence | PackedSequence
+
+
+class PaddedInput(NamedTuple):
+    x: Float[Tensor, "b n d"]
+    mask: Bool[Tensor, "b n"]
+    abs_positions: Float[Tensor, "b n c"]  # c=1 for 1-D (text), c=2 for grids, c=3 for 3-D
+
+
+class PackedInput(NamedTuple):
+    x: Float[Tensor, "nt d"]
+    cu_seqlens: Int[Tensor, "bp1"]
+    max_seqlen: int
+    abs_positions: Float[Tensor, "nt c"]  # c matches PaddedInput convention
+
+
+SequenceInput = PaddedInput | PackedInput
+
+
+def to_seq_info(inp: SequenceInput) -> SequenceInfo:
+    match inp:
+        case PaddedInput(mask=mask):
+            return PaddedSequence(mask=mask)
+        case PackedInput(cu_seqlens=cu, max_seqlen=ms):
+            return PackedSequence(cu_seqlens=cu, max_seqlen=ms)
+
+
+def make_padded_input(x: Tensor, mask: Bool[Tensor, "b n"]) -> PaddedInput:
+    """Build PaddedInput with sequential 1-D positions (shape b n 1)."""
+    n = x.shape[1]
+    pos = torch.arange(n, device=x.device, dtype=x.dtype)
+    positions = pos.unsqueeze(0).unsqueeze(-1).expand(x.shape[0], -1, -1)  # b n 1
+    return PaddedInput(x=x, mask=mask, abs_positions=positions)
+
+
+def make_packed_input(x: Tensor, cu_seqlens: Int[Tensor, "bp1"], max_seqlen: int) -> PackedInput:
+    """Build PackedInput with per-token sequential 1-D positions (shape nt 1).
+
+    Delegates to :func:`position_ids_from_packed`, which is ``torch.export``-compatible.
+    """
+    seq = PackedSequence(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+    positions = position_ids_from_packed(seq).to(x.dtype).unsqueeze(-1)  # nt 1
+    return PackedInput(x=x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, abs_positions=positions)
+
+
+def padded_to_key_padding_mask(seq: PaddedSequence) -> Bool[Tensor, "b n"]:
+    """Return True where tokens are valid (mirrors seq.mask)."""
+    return seq.mask
+
+
+def packed_batch_size(seq: PackedSequence) -> int:
+    return int(seq.cu_seqlens.shape[0]) - 1
+
+
+def make_padded(mask: Bool[Tensor, "b n"]) -> PaddedSequence:
+    return PaddedSequence(mask=mask)
+
+
+def make_packed(cu_seqlens: Int[Tensor, "bp1"], max_seqlen: int) -> PackedSequence:
+    return PackedSequence(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
+
+
+def padded_to_packed(inp: PaddedInput) -> PackedInput:
+    """Remove padding tokens from a padded batch, producing a flat packed tensor.
+
+    Output shape is data-dependent (nt = mask.sum()); not compatible with torch.export.
+    """
+    lengths = inp.mask.sum(dim=1)
+    cu = lengths_to_cu_seqlens(lengths)
+    max_seqlen = int(lengths.max().item())
+    return PackedInput(
+        x=inp.x[inp.mask],
+        cu_seqlens=cu,
+        max_seqlen=max_seqlen,
+        abs_positions=inp.abs_positions[inp.mask],
+    )
+
+
+def packed_to_padded(inp: PackedInput) -> PaddedInput:
+    """Re-pad a packed tensor to (b, max_seqlen, d), filling unused positions with zeros.
+
+    Uses scatter — compatible with torch.export and torch.compile.
+    """
+    cu = inp.cu_seqlens
+    b = cu.shape[0] - 1
+    n = inp.max_seqlen
+    # Build a (nt,) tensor of batch indices and a (nt,) tensor of within-sequence positions.
+    lengths = cu[1:] - cu[:-1]  # (b,)
+    batch_idx = torch.repeat_interleave(
+        torch.arange(b, device=cu.device, dtype=torch.long), lengths
+    )  # (nt,)
+    pos_idx = position_ids_from_packed(PackedSequence(cu_seqlens=cu, max_seqlen=n))  # (nt,)
+    d, c = inp.x.shape[-1], inp.abs_positions.shape[-1]
+    x_out = inp.x.new_zeros(b, n, d)
+    x_out[batch_idx, pos_idx] = inp.x
+    pos_out = inp.abs_positions.new_zeros(b, n, c)
+    pos_out[batch_idx, pos_idx] = inp.abs_positions
+    mask = torch.zeros(b, n, dtype=torch.bool, device=inp.x.device)
+    mask[batch_idx, pos_idx] = True
+    return PaddedInput(x=x_out, mask=mask, abs_positions=pos_out)
+
+
+def lengths_to_cu_seqlens(lengths: Int[Tensor, "b"]) -> Int[Tensor, "bp1"]:
+    zero = torch.zeros(1, dtype=lengths.dtype, device=lengths.device)
+    return torch.cat([zero, lengths.cumsum(0)])
+
+
+def position_ids_from_packed(seq: PackedSequence) -> Int[Tensor, "nt"]:
+    """Build per-token position indices [0..len_i-1] for each document in the pack.
+
+    The result at flat index i is ``i - cu_seqlens[batch_of_i]``, i.e. the within-document
+    position of that token.
+
+    Compatible with ``torch.export`` and ``torch.compile``: no Python-side iteration over
+    data-dependent lengths.  The key identity is::
+
+        pos_idx = arange(nt) - cu[batch_idx]
+
+    where ``batch_idx`` is built via ``repeat_interleave``, which is a supported dynamic op.
+    """
+    cu = seq.cu_seqlens
+    b = cu.shape[0] - 1
+    lengths = cu[1:] - cu[:-1]  # (b,) — kept as a tensor, never converted to Python list
+    batch_idx = torch.repeat_interleave(
+        torch.arange(b, device=cu.device, dtype=torch.long), lengths
+    )  # (nt,)
+    nt = batch_idx.shape[0]
+    return torch.arange(nt, device=cu.device, dtype=torch.long) - cu[batch_idx]

stackformers-3.8.0.dist-info/METADATA

@@ -0,0 +1,204 @@
+Metadata-Version: 2.4
+Name: stackformers
+Version: 3.8.0
+Summary: Typed, composable, SOLID transformer library for PyTorch
+Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
+Project-URL: Repository, https://github.com/Red-Eyed/stackformers
+Project-URL: Bug Tracker, https://github.com/Red-Eyed/stackformers/issues
+Author-email: Vadym Stupakov <vadim.stupakov@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2026 Vadym Stupakov
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: attention,deep-learning,pytorch,transformer
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: beartype>=0.18
+Requires-Dist: einops>=0.7
+Requires-Dist: jaxtyping>=0.2
+Requires-Dist: numpy>=2.4.4
+Requires-Dist: pydantic>=2.0
+Requires-Dist: torch>=2.2
+Description-Content-Type: text/markdown
+
+# stackformers
+
+Typed, composable transformer library for PyTorch.
+Every architectural choice — positional encoding, normalization, feedforward variant — is an injected dependency, not a constructor flag.
+
+```bash
+uv add stackformers
+```
+
+---
+
+## Why
+
+Most transformer libraries grow into a tangle of `if self.use_rope`, `if self.window_size is not None`, and god-config objects with thirty nullable fields. Adding a new variant means touching existing code.
+
+stackformers takes a different approach:
+
+- **Swap any component without touching anything else** — `SelfAttention(config, pos_encoding=RoPE)` vs `SelfAttention(config, pos_encoding=ALiBi)` — same call site, different object
+- **No `None` checks in `forward()`** — `NoPosEncoding` is a real object that passes q/k unchanged; the branch never exists
+- **Sealed sequence unions** — `PaddedInput | PackedInput` instead of optional `cu_seqlens` and `mask` arguments that conflict with each other
+- **`torch.compile` / `torch.export` safe** — no Python control flow on tensors inside any `forward()`
+- **Structural protocols** — bring your own implementation; no ABC inheritance required
+
+---
+
+## Quick start
+
+### Zero boilerplate
+
+```python
+import torch
+from stackformers import TransformerEncoder, plain_encoder_config, make_padded_input
+
+model = TransformerEncoder(plain_encoder_config(dim=512, heads=8, num_layers=6))
+
+x    = torch.randn(2, 128, 512)
+mask = torch.ones(2, 128, dtype=torch.bool)
+out  = model(make_padded_input(x, mask))   # (2, 128, 512)
+```
+
+Switch to packed (variable-length, no padding waste) — same weights:
+
+```python
+from stackformers import make_packed_input
+
+cu  = torch.tensor([0, 64, 128], dtype=torch.int32)
+out = model(make_packed_input(x_flat, cu, max_seqlen=64))  # (128, 512)
+```
+
+Causal LM backbone:
+
+```python
+plain_encoder_config(dim=768, heads=12, num_layers=12, causal=True)
+```
+
+Sliding-window local attention (O(n · w)):
+
+```python
+from stackformers import windowed_encoder_config
+windowed_encoder_config(dim=512, heads=8, num_layers=6, window_size=128)
+```
+
+Encoder–decoder:
+
+```python
+from stackformers import TransformerDecoder, plain_decoder_config
+
+model = TransformerDecoder(plain_decoder_config(dim=512, heads=8, num_layers=6))
+out   = model(make_padded_input(x, mask), make_padded_input(context, ctx_mask))
+```
+
+### Explicit config
+
+Full control with JSON round-trip via `kind` discriminators:
+
+```python
+from stackformers import (
+    TransformerEncoderConfig, TransformerEncoder,
+    SelfAttentionConfig, SwiGLUConfig, RMSNormConfig, RoPE1DConfig,
+    make_padded_input,
+)
+
+cfg = TransformerEncoderConfig(
+    attn=SelfAttentionConfig(dim=512, heads=8, dim_head=64, causal=False),
+    ff=SwiGLUConfig(dim=512, mult=4.0),
+    norm=RMSNormConfig(dim=512),
+    pos_encoding=RoPE1DConfig(dim_head=64),
+    num_layers=6,
+)
+model = TransformerEncoder(cfg)
+
+# Serialise / restore
+cfg2 = TransformerEncoderConfig.model_validate(cfg.model_dump())
+```
+
+### Custom wiring
+
+Wire layers yourself when presets aren't enough:
+
+```python
+from stackformers import (
+    SelfAttention, SwiGLU, TransformerLayer, Encoder, RMSNorm,
+    RotaryEmbedding1D,
+    SelfAttentionConfig, SwiGLUConfig, RMSNormConfig, RoPE1DConfig,
+)
+
+pos  = RotaryEmbedding1D(RoPE1DConfig(dim_head=64))
+attn = SelfAttention(SelfAttentionConfig(dim=512, heads=8, dim_head=64), pos_encoding=pos)
+
+layers = [
+    TransformerLayer(
+        self_attn=attn,
+        ff=SwiGLU(SwiGLUConfig(dim=512)),
+        norm_attn=RMSNorm(RMSNormConfig(dim=512)),
+        norm_ff=RMSNorm(RMSNormConfig(dim=512)),
+    )
+    for _ in range(6)
+]
+encoder = Encoder(layers=layers, final_norm=RMSNorm(RMSNormConfig(dim=512)))
+```
+
+---
+
+## What's included
+
+| Area | Variants |
+|------|----------|
+| Self-attention | Global, sliding-window (local); padded and packed backends; GQA / MQA |
+| Cross-attention | Global; padded and packed backends |
+| Positional encoding | RoPE-1D, RoPE-2D, none (null object) |
+| Feedforward | SwiGLU, GEGLU |
+| Normalization | RMSNorm, LayerNorm |
+| Presets | Encoder, Decoder, CrossAttender |
+
+On CUDA with fp16/bf16 the packed path uses `torch.nn.attention.varlen.varlen_attn`. CPU and fp32 fall back to a scatter-to-padded SDPA — correct everywhere, fast where it matters.
+
+---
+
+## Development
+
+```bash
+git clone <repo> && cd stackformers
+uv sync --group dev
+
+just fmt      # format
+just lint     # lint
+just types    # type-check
+just test     # test
+just check    # full CI gate
+```
+
+---
+
+## License
+
+See [LICENSE](LICENSE).

stackformers-3.8.0.dist-info/RECORD

@@ -0,0 +1,46 @@
+stackformers/__init__.py,sha256=kXFjn1AbcqhP3Pf2TynXdHkaeB3qivOSsGDU5rrtgWA,3787
+stackformers/config.py,sha256=tPRmrwvO7t1rGw2KA2pna-Tt-5He4QaBYh1uF464t0I,704
+stackformers/cross_attender.py,sha256=VFrDDLWv_36INFUiScYl8zc4ujeS7w5tP0v7ksu7kNo,1506
+stackformers/decoder.py,sha256=PsjssyIEM6tqUWasQ1AbqLj-pfTHjVG2nnQy8skYf4E,1824
+stackformers/encoder.py,sha256=GBe3HxZUHCdpfa8WuQZMXM-yJR1FdIF48HZbdg5gpGI,701
+stackformers/layers.py,sha256=uM_85QhaGvlwXpkmC1IDZRtnPIr-DyVKcLz38626EdE,941
+stackformers/sequence.py,sha256=8THo2FrN6InTEo1f9fSGlxMsP2wLcF3oVn3PQHAIa5w,5788
+stackformers/attention/README.md,sha256=7UGXLTcJ9V3Imj9wgpzbaDHUKnfOXS4qU3zZ5oD_QQo,2544
+stackformers/attention/__init__.py,sha256=2t91tmmlRy8Ww3gx0B-QX4_9LsQh72yl6smXyrxIJX4,510
+stackformers/attention/bias.py,sha256=zIb8iEdqdO1gvnVesWJwTP60qAv1I3elldv97chkHpM,298
+stackformers/attention/config.py,sha256=9QRBQ3akBXAHIHNDa-8_uHfAEH-etcaKPcc0RTRo1fM,5039
+stackformers/attention/cross_attn.py,sha256=6Ak_DHlsCZu493YqwiTd8sYjMM52h-qNw39u9Rydujg,3965
+stackformers/attention/ops.py,sha256=_V82bYaDlKurlOX5hL4nfObqRXuU2BiLs5zqB3Ij8es,6097
+stackformers/attention/protocols.py,sha256=xl3ZkXUM93poT_2QxPGlRdS_7z6tf_1yNT1gjnSY9-I,1167
+stackformers/attention/self_attn.py,sha256=6km3lFhkAiFctqDkI_voDXqGHEV2wyKssworfq44p3k,3933
+stackformers/feedforward/README.md,sha256=rgOSpWmBRmoApaai1ZpRjiYtY1auCbqciP-lTkPwJkQ,503
+stackformers/feedforward/__init__.py,sha256=AviF8Ib76eU5J3zuuiz4A4u1bl8qoB-UZ0LCS_Xgegk,212
+stackformers/feedforward/config.py,sha256=XEAev9SDKmZzHFjFPfE3WvRpWYcXMK_9Sy7fmWd0OU0,2103
+stackformers/feedforward/factory.py,sha256=FiLaR1SZ-CoU4uGZQaJizii718kl3kw0GsMmVrDuL5E,761
+stackformers/feedforward/geglu.py,sha256=lEAZHyh4xfIxWZzTsxffr49lyQphXDpkqCQIs77Y7Lg,1200
+stackformers/feedforward/protocols.py,sha256=ZxRJDYjGAyu7ruxWcBEblds7CMZQJK6kNm_qxNN4Zbg,380
+stackformers/feedforward/relu_squared.py,sha256=8rhVGUaoXMah_prkz9SkTVNsHJ6AxpNjda2w9OPugbo,953
+stackformers/feedforward/swiglu.py,sha256=87b0kH2sRWB4U2wpkGGt5fmOLFrHoM_0jkhYKhQRZvY,1114
+stackformers/norm/README.md,sha256=NMENfZoBWcaVxcth4SBVn2RZD7cyL90PfcTxDrz_f0I,495
+stackformers/norm/__init__.py,sha256=2ZbOQJbXg8SNxVo5AsmI1rsEMWw33cu6sQI-9DM67YQ,116
+stackformers/norm/config.py,sha256=HIbn8NLcdAn_WIrSMaCwe0kMUi0Fh8Ykapz6E-P_RsQ,469
+stackformers/norm/factory.py,sha256=AoaKgpn1aHGYNb4dqqGQiCd5d716agVroSUoJ4XLMQc,559
+stackformers/norm/protocols.py,sha256=WWhqBqBFG5TSZe4U42tzDhT0X8nWJaRYDU9OnQ8uem4,355
+stackformers/positional/README.md,sha256=R-QYDL7PiSvd7NwonEW7Mwcv7WtEHYzTNE5fIKHqJ8s,1114
+stackformers/positional/__init__.py,sha256=1bqfSD6ekk35lQtm6tiEiwnvxeu3I6Gno1Um1jTr8oM,633
+stackformers/positional/config.py,sha256=gT5w0BJnX5lvZRIav3YwkSvkbfDt6it2oq782OGuLpA,2141
+stackformers/positional/factory.py,sha256=ILT8yZUEvvEuM6tkHyXSwc6OAQ37HkHIsrpUOeZJ7Rs,992
+stackformers/positional/learned.py,sha256=nmXGv7HnNgadcq7LZ5OrmPlKT1qcPvAVKuQR53Ow55U,1855
+stackformers/positional/none.py,sha256=lq4WREpY572jux3_V-_-9xG0ZvCgJUddMQTMfBZBwVs,993
+stackformers/positional/protocols.py,sha256=uLFHYJrHfYedBxzxe3-2Apuq4VaY0zMgbN1gT8Uc0Uw,1014
+stackformers/positional/rope1d.py,sha256=Y2Dn5DjO14lcQ7PmjQ4zZLxA4tefuB3yBMZz9x5QgbY,3851
+stackformers/positional/rope2d.py,sha256=4VXQ0tDUnlFR386LlyWfOx9U85XwvH7_qVPJ-8mIVas,2319
+stackformers/presets/README.md,sha256=uWaW8dGlGth2YNaE8TTWbaQTbLnojnp4g59Ez9lONao,2628
+stackformers/presets/__init__.py,sha256=vlQagWfVRHFnH6GgqA8EUe_rEx5jw9eoif2SSyPrIxM,1024
+stackformers/presets/cross_attender.py,sha256=yZG4oEnZqTUJwXuPZKF-fRdtdWpJ7dAO6WaECIFMQpY,3838
+stackformers/presets/decoder.py,sha256=SChy1PlIUQ82zGeXcvJyPae60nbvfCC1_pf7OBb9ngI,4699
+stackformers/presets/encoder.py,sha256=Hj6nDkSUPxeXMdBsBh3hl41JlNkHtZzWW7W05JuHLkA,4990
+stackformers-3.8.0.dist-info/METADATA,sha256=WmSQC9jHW1XJFTgOD1ovm0nnlby_7GV0n-R1-uoGx3I,6846
+stackformers-3.8.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+stackformers-3.8.0.dist-info/licenses/LICENSE,sha256=6rViNgu6RIlBLiKHptpp9oPrAZBbLJlw_1IPUNGvMOg,1071
+stackformers-3.8.0.dist-info/RECORD,,

stackformers-3.8.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any