stackformers

Typed, composable transformer library for PyTorch. Every architectural choice — positional encoding, normalization, feedforward variant — is an injected dependency, not a constructor flag.

uv add stackformers

Runnable examples cover variable-length layouts, causal language modeling, image patches, continuous point coordinates, user subclassing, and cached ONNX Runtime deployment. See the examples guide.

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

Import the public API with import stackformers as sf. Configs, models, protocols, base classes, factories, sequence helpers, and projection conversions are available directly through sf.

Zero boilerplate

import torch
import stackformers as sf

model = sf.TransformerEncoder(sf.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(sf.make_padded_input(x, mask))   # (2, 128, 512)

Switch to packed (variable-length, no padding waste) — same weights:

cu  = torch.tensor([0, 64, 128], dtype=torch.int32)
out = model(sf.make_packed_input(x_flat, cu, max_seqlen=64))  # (128, 512)

Causal LM backbone:

sf.plain_encoder_config(dim=768, heads=12, num_layers=12, causal=True)

Sliding-window local attention (O(n · w)):

sf.windowed_encoder_config(dim=512, heads=8, num_layers=6, window_size=128)

Variable-width encoder — one model and head dimension per Transformer block:

config = sf.variable_width_encoder_config(
    d_models=[512, 512, 768, 768, 1024],
    dim_heads=[64, 64, 64, 64, 128],
)
model = sf.VariableWidthTransformerEncoder(config)

Omitting heads derives each block's head count as d_model // dim_head. Set heads=12 or heads=[12, 12, 16, 16, 16] to choose the internal attention width independently of d_models. dim_heads, heads, causal, ff_mult, dropout, and norm_placement accept either a shared value or a list with one entry per block. The preset inserts a bias-free learned projection before a block when its residual width differs from the preceding block; equal-width neighbors use an identity. The convenience factory expands each width pair into a VariableWidthEncoderLayerConfig whose attention, feed-forward, norm, positional encoding, and attention bias are all explicit. Construct these layer configs directly to replace any component; the model does not rebuild or override them. The input width must match the first layer's attn.dim, and the output uses the last layer's width.

Reference: Wu et al., “Variable-Width Transformers” (2026). The paper uses parameter-free residual resizing; this preset instead makes each width transition a learned linear projection, so it is a related experimental variant rather than an exact replica.

Encoder–decoder:

model = sf.TransformerDecoder(sf.plain_decoder_config(dim=512, heads=8, num_layers=6))
out   = model(sf.make_padded_input(x, mask), sf.make_padded_input(context, ctx_mask))

For autoregressive export, build cross K/V once and pass both required cache tensors to every one-token decoder invocation:

import torch

import stackformers as sf

model.eval()
cache_builder = sf.DecoderCrossAttentionCacheBuilder(model).eval()
cached_decoder = sf.CachedDecoderWrapper(model).eval()

cross_cache = cache_builder(context_input)
self_kv_cache = torch.empty(num_layers, 2, batch, kv_heads, 0, dim_head)
step_i = torch.tensor([0], dtype=torch.int64)
cache_onnx = torch.onnx.export(cache_builder, (context_input,), dynamo=True)
decoder_onnx = torch.onnx.export(
    cached_decoder,
    (target_token, cross_cache, self_kv_cache, step_i),
    dynamo=True,
)

The cache builder runs once per encoder context. Each decoder call consumes exactly one target token, the immutable cross cache, the growing self cache, and step_i; it returns the output token and a self cache extended by one position. K/V tensors retain kv_heads rather than expanded query heads. Cache construction positions cross K once, while cached decoding positions each new self K and both self/cross Q, so all built-in positional encodings remain available. Batch, source, and past-target axes are dynamic during export; model width, layer count, head geometry, and the one-token axis remain static because weights constrain them.

See the complete eager, ONNX export, and raw ONNX Runtime demonstration in examples/encoder_decoder_onnxruntime.py.

Explicit config

Full control with JSON round-trip via kind discriminators:

import stackformers as sf

cfg = sf.TransformerEncoderConfig(
    attn=sf.SelfAttentionConfig(dim=512, heads=8, dim_head=64, causal=False),
    ff=sf.SwiGLUConfig(dim=512, mult=4.0),
    norm=sf.RMSNormConfig(dim=512),
    pos_encoding=sf.RoPE1DConfig(dim_head=64),
    num_layers=6,
)
model = sf.TransformerEncoder(cfg)

# Serialise / restore
cfg2 = sf.TransformerEncoderConfig.model_validate(cfg.model_dump())

norm_placement is available on TransformerEncoderConfig, TransformerDecoderConfig, and CrossAttenderConfig, as well as their plain-config helpers. It accepts four layouts and defaults to "pre", so existing constructors, serialized configs, checkpoints, outputs, and gradients retain the previous behavior when the field is omitted. For a residual branch F with norm N:

Sandwich placement creates independent pre- and post-branch norms. Reordered placement follows the OLMo 2 residual layout; enable QK-Norm separately in the attention config when reproducing the broader OLMo 2 stabilization recipe. Preset construction maps each value to a focused encoder, decoder, or cross-attender layer class. Decoder and cross-attender placement applies to the target/query residual stream; the context sequence is not normalized or mutated by these layers. Ordinary and cached decoder layers call the same normalization/residual equation functions. Attention execution is supplied through small call adapters, while norm and projection modules remain registered under their existing checkpoint names. Internal context variants encode either omission or a complete context value; cached contexts keep K/V and their validity mask together.

Cached deployment adapters support exact built-in attention, decoder-stack, and decoder-layer classes. Unsupported subclasses raise TypeError because the adapters cannot infer overridden forward behavior. Custom collaborators continue to work in ordinary decoding. Direct cached-layer calls require both cross cache and context, or neither; a half-present pair raises ValueError.

Custom wiring

Wire layers yourself when presets aren't enough:

import stackformers as sf

pos  = sf.RotaryEmbedding1D(sf.RoPE1DConfig(dim_head=64))
attn = sf.SelfAttention(sf.SelfAttentionConfig(dim=512, heads=8, dim_head=64), pos_encoding=pos)

layers = [
    sf.TransformerLayer(
        self_attn=attn,
        ff=sf.SwiGLU(sf.SwiGLUConfig(dim=512)),
        norm_attn=sf.build_norm(sf.RMSNormConfig(dim=512)),
        norm_ff=sf.build_norm(sf.RMSNormConfig(dim=512)),
    )
    for _ in range(6)
]
encoder = sf.Encoder(layers=layers, final_norm=sf.build_norm(sf.RMSNormConfig(dim=512)))

Convert projection layouts

Load an ordinary checkpoint, then recursively pack supported self-attention, cross-attention, SwiGLU, GEGLU, and HardSwishGLU layers:

import stackformers as sf

model.load_state_dict(checkpoint)
model = sf.to_packed(model)
# Run or save the packed model.
model = sf.to_unpacked(model)

These helpers replace registered children in place and return the root, which may itself be replaced. Repeated calls keep layers already in the requested layout. Shared layer references and retained child modules stay shared. Built-in preset configs follow the new layout; save model.config.model_dump_json() alongside its state dict to reconstruct that layout later. Caller-owned config objects are left unchanged. Custom subclasses keep their own behavior and metadata; supported children are still visited.

Weights round-trip exactly; outputs match within dtype tolerance because matrix multiplication rounding can change. Convert before constructing an optimizer: replaced projection parameters do not migrate optimizer state, existing gradients, hooks, or cross-layer parameter aliases. Groups with mixed requires_grad fail before mutation. Cache wrappers support both ordinary and packed attention with the same calls and cache tensor layouts. Pack the model before constructing cache wrappers, or convert a containing module that registers the model and its existing wrappers together so their attention references remain shared. Packing a wrapper alone converts only its registered children; references outside that tree still point to the original attention modules. Projection packing is independent of the PackedInput sequence layout.

What's included

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

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

The quality gate checks the library, tests, and examples against Python 3.11. Ruff enforces annotations, common bug patterns, typing imports, and qualified, current suppressions. Pyrefly uses the strict preset with explicit/implicit Any, untyped returns, missing annotations, and stale suppressions treated as errors. Tool minimums are recorded in both the development dependencies and checker configs.

Module classes expose __call__ = forward only under TYPE_CHECKING. This gives callers the declared input/output contract while leaving PyTorch's hooks and runtime dispatch intact. Local Tensor annotations at built-in PyTorch projection boundaries rely on those modules' tensor-return contracts; they do not validate tensor shapes. Experimental kernel adapters capture the vendor's declared return type and validate the actual response before exposing Result[Tensor, VarlenFallback]. Export shape trees use opaque object leaves because PyTorch validates their module-dependent structure.

Core admission checks return Result[T, E] from stackformers.result, the closed union Ok[T] | Err[E]. Both variants are generic NamedTuples with a value or error and a literal kind tag. Match Ok(value=value) and Err(error=error) before accessing their payloads; keyword class patterns work with strict PyTorch export and full-graph compilation. No external Result package or unchecked unwrap method is required.

Recoverable core failures remain error data. Vendor adapters catch expected third-party exceptions and preserve their original causes; cache admission validates before construction without raising and catching its own expected failures. Existing public adapters retain tensor/configuration returns and raise when rejected inputs prevent that operation from proceeding. Pydantic still presents configuration failures as ValidationError; the experimental backend retains tensor-or-None returns and fallback warnings. Unexpected programming and PyTorch failures can propagate. See Result records.

License

See LICENSE.