stackformers 3.8.0__py3-none-any.whl

stackformers/__init__.py

@@ -0,0 +1,136 @@
+"""stackformers public API."""
+
+from importlib.metadata import version
+
+__version__ = version("stackformers")
+
+from stackformers.attention.config import CrossAttentionConfig, SelfAttentionConfig
+from stackformers.attention.cross_attn import CrossAttention
+from stackformers.attention.protocols import CrossAttn, SelfAttn
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.config import DecoderConfig, EncoderConfig, LayerConfig
+from stackformers.cross_attender import CrossAttenderLayer, CrossAttenderStack
+from stackformers.decoder import Decoder, DecoderLayer
+from stackformers.encoder import Encoder
+from stackformers.feedforward.config import FeedForwardConfig
+from stackformers.feedforward.factory import build_ff
+from stackformers.feedforward.protocols import FeedForward
+from stackformers.feedforward.swiglu import SwiGLU
+from stackformers.layers import TransformerLayer
+from stackformers.norm.config import LayerNormConfig, RMSNormConfig
+from stackformers.norm.factory import NormConfig, build_norm
+from stackformers.norm.protocols import Norm
+from stackformers.positional.config import (
+    NoPosEncodingConfig,
+    PosEncodingConfig,
+    RoPE1DConfig,
+    RoPE2DConfig,
+)
+from stackformers.positional.factory import build_pos_encoding
+from stackformers.positional.none import NoPosEncoding
+from stackformers.positional.protocols import PosEncoding
+from stackformers.positional.rope1d import RotaryEmbedding1D
+from stackformers.positional.rope2d import RotaryEmbedding2D
+from stackformers.presets.cross_attender import (
+    CrossAttender,
+    CrossAttenderConfig,
+    plain_cross_attender_config,
+)
+from stackformers.presets.decoder import (
+    TransformerDecoder,
+    TransformerDecoderConfig,
+    plain_decoder_config,
+)
+from stackformers.presets.encoder import (
+    TransformerEncoder,
+    TransformerEncoderConfig,
+    plain_encoder_config,
+    windowed_encoder_config,
+)
+from stackformers.sequence import (
+    PackedInput,
+    PackedSequence,
+    PaddedInput,
+    PaddedSequence,
+    SequenceInfo,
+    SequenceInput,
+    lengths_to_cu_seqlens,
+    make_packed,
+    make_packed_input,
+    make_padded,
+    make_padded_input,
+    position_ids_from_packed,
+    to_seq_info,
+)
+
+__all__ = [
+    "__version__",
+    # sequences
+    "PaddedSequence",
+    "PackedSequence",
+    "SequenceInfo",
+    "PaddedInput",
+    "PackedInput",
+    "SequenceInput",
+    "make_padded",
+    "make_packed",
+    "make_padded_input",
+    "make_packed_input",
+    "to_seq_info",
+    "lengths_to_cu_seqlens",
+    "position_ids_from_packed",
+    # protocols
+    "PosEncoding",
+    "SelfAttn",
+    "CrossAttn",
+    "FeedForward",
+    "Norm",
+    # configs — attention
+    "SelfAttentionConfig",
+    "CrossAttentionConfig",
+    # configs — ff / layer / encoder / decoder
+    "FeedForwardConfig",
+    "LayerConfig",
+    "EncoderConfig",
+    "DecoderConfig",
+    # configs — norm
+    "RMSNormConfig",
+    "LayerNormConfig",
+    "NormConfig",
+    # configs — positional
+    "RoPE1DConfig",
+    "RoPE2DConfig",
+    "NoPosEncodingConfig",
+    "PosEncodingConfig",
+    # positional
+    "NoPosEncoding",
+    "RotaryEmbedding1D",
+    "RotaryEmbedding2D",
+    # attention
+    "SelfAttention",
+    "CrossAttention",
+    # feedforward
+    "SwiGLU",
+    # builders
+    "build_norm",
+    "build_ff",
+    "build_pos_encoding",
+    # transformer blocks
+    "TransformerLayer",
+    "Encoder",
+    "DecoderLayer",
+    "Decoder",
+    "CrossAttenderLayer",
+    "CrossAttenderStack",
+    # presets
+    "TransformerEncoderConfig",
+    "TransformerEncoder",
+    "plain_encoder_config",
+    "windowed_encoder_config",
+    "TransformerDecoderConfig",
+    "TransformerDecoder",
+    "plain_decoder_config",
+    "CrossAttenderConfig",
+    "CrossAttender",
+    "plain_cross_attender_config",
+]

stackformers/attention/README.md

@@ -0,0 +1,21 @@
+# attention
+
+Multi-head self-attention and cross-attention.
+
+## Design
+
+**Two configs, two modules.** `SelfAttentionConfig` and `CrossAttentionConfig` are separate types because their parameter sets differ: self-attention has `causal` and `window_size`; cross-attention has neither.
+
+**Static vs runtime dispatch.** The attention *type* (global vs sliding-window) is determined at construction time via `window_size: int | None` in `SelfAttentionConfig`. The *backend* (padded vs packed) is determined at runtime: `SelfAttention.forward` and `CrossAttention.forward` dispatch on `SequenceInput` via `match`. This means one model handles both training (packed) and inference (padded) without swapping classes.
+
+**Kernel ops in `ops.py`.** All low-level attention math lives in `attention/ops.py`: mask builders, `packed_attn` (thin `varlen_attn` wrapper), `padded_sdpa`, and `packed_attn_or_fallback`. The modules stay thin dispatchers.
+
+**Packed attention with automatic fallback.** `PackedInput` normally routes to `varlen_attn`, which requires CUDA and `float16`/`bfloat16`. When those conditions are not met (CPU, float32, `torch.export`), `packed_attn_or_fallback` transparently scatters q/k/v to padded layout, runs `F.scaled_dot_product_attention`, and gathers the result back to packed — same weights, same output shape `(nt, h, d)`. A `UserWarning` is emitted so callers know they are on the slow path. Dropout works on the fallback path; it is silently skipped only when `varlen_attn` is used (which does not accept a dropout argument).
+
+**`torch.export` and `torch.compile` compatibility.** All scatter/gather helpers in `ops.py` (`_cu_to_indices`, `_packed_heads_to_padded`, `_padded_heads_to_packed`) are fully export-compatible. The key design constraint is that per-document length iteration must never use Python-side loops over data-dependent values (i.e., no `.tolist()` followed by a list comprehension). Instead, `_cu_to_indices` uses `torch.repeat_interleave` to produce `batch_idx` and derives `pos_idx` via the identity `arange(nt) - cu[batch_idx]`. This keeps the computation graph purely symbolic even when `nt` (total tokens) is a data-dependent unbacked symbol.
+
+**GQA / MQA.** Set `kv_heads < heads` in either config. Key and value heads are repeated via `einops.repeat` before the dot-product — no attention-specific changes needed.
+
+## Extending
+
+To add a new attention variant, create a new `nn.Module` that satisfies the `SelfAttn` or `CrossAttn` protocol. No kernel registration or factory changes required.

stackformers/attention/__init__.py

@@ -0,0 +1,16 @@
+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, CrossAttn, SelfAttn
+from stackformers.attention.self_attn import SelfAttention
+
+__all__ = [
+    "SelfAttentionConfig",
+    "CrossAttentionConfig",
+    "SelfAttn",
+    "CrossAttn",
+    "AttnBias",
+    "SelfAttention",
+    "CrossAttention",
+    "NoAttnBias",
+]

stackformers/attention/bias.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+import torch.nn as nn
+
+from stackformers.sequence import PaddedInput
+
+
+class NoAttnBias(nn.Module):
+    """Null object for AttnBias — signals no bias so varlen_attn remains available."""
+
+    def __call__(self, input: PaddedInput) -> None:
+        return None

stackformers/attention/config.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import warnings
+
+from pydantic import BaseModel, Field, model_validator
+
+_ALIGN = 64  # tensor-core alignment for FP16/BF16
+
+
+def _validate_attn_dims(dim: int, heads: int, dim_head: int, kv_heads: int | None) -> None:
+    if dim % heads != 0:
+        raise ValueError(f"dim ({dim}) must be divisible by heads ({heads})")
+    kv_h = kv_heads if kv_heads is not None else heads
+    if heads % kv_h != 0:
+        raise ValueError(f"heads ({heads}) must be divisible by kv_heads ({kv_h})")
+    if dim_head % _ALIGN != 0:
+        warnings.warn(
+            f"dim_head={dim_head} is not a multiple of {_ALIGN}. "
+            "Unaligned head dimension reduces GPU throughput on tensor-core hardware. "
+            f"Nearest aligned values: {(dim_head // _ALIGN) * _ALIGN} or "
+            f"{(dim_head // _ALIGN + 1) * _ALIGN}.",
+            UserWarning,
+            stacklevel=3,
+        )
+    projected = heads * dim_head
+    if projected != dim:
+        warnings.warn(
+            f"heads * dim_head = {heads} * {dim_head} = {projected}, "
+            f"which does not equal dim={dim}. "
+            "The Q/K/V projections are non-square — this is valid but may not be intentional. "
+            f"Standard transformers use dim_head = dim // heads = {dim // heads}.",
+            UserWarning,
+            stacklevel=3,
+        )
+
+
+class SelfAttentionConfig(BaseModel):
+    dim: int = Field(
+        gt=0,
+        description="Model (embedding) dimension — input and output width of the attention sublayer.",  # noqa: E501
+    )
+    heads: int = Field(default=8, gt=0, description="Number of query heads.")
+    dim_head: int = Field(
+        default=64,
+        gt=0,
+        description=(
+            "Dimension per head. Output projection width is heads * dim_head,"
+            " which need not equal dim (non-square projections are valid)."
+        ),
+    )
+    kv_heads: int | None = Field(
+        default=None,
+        description=(
+            "Number of key/value heads. None = equal to heads (standard MHA)."
+            " Set to 1 for MQA or any divisor of heads for GQA."
+        ),
+    )
+    causal: bool = Field(
+        default=False,
+        description=(
+            "Apply a causal mask so each position attends only to itself and earlier positions."
+        ),
+    )
+    dropout: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Dropout probability applied to the output projection, during training only.",
+    )
+    window_size: int | None = Field(
+        default=None,
+        description=(
+            "Sliding-window width. None = global attention."
+            " A positive integer restricts each token to a local window of that width."
+        ),
+    )
+    qk_norm: bool = Field(
+        default=False,
+        description=(
+            "Apply RMSNorm to queries and keys (in head space) before the dot product."
+            " Stabilises logit scale at large depth or large dim_head."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def _validate(self) -> "SelfAttentionConfig":
+        _validate_attn_dims(self.dim, self.heads, self.dim_head, self.kv_heads)
+        return self
+
+    @property
+    def effective_kv_heads(self) -> int:
+        return self.kv_heads if self.kv_heads is not None else self.heads
+
+    @property
+    def groups(self) -> int:
+        return self.heads // self.effective_kv_heads
+
+
+class CrossAttentionConfig(BaseModel):
+    dim: int = Field(
+        gt=0,
+        description="Model (embedding) dimension — input and output width of the attention sublayer.",  # noqa: E501
+    )
+    heads: int = Field(default=8, gt=0, description="Number of query heads.")
+    dim_head: int = Field(
+        default=64,
+        gt=0,
+        description=(
+            "Dimension per head. Output projection width is heads * dim_head,"
+            " which need not equal dim (non-square projections are valid)."
+        ),
+    )
+    kv_heads: int | None = Field(
+        default=None,
+        description=(
+            "Number of key/value heads. None = equal to heads (standard MHA)."
+            " Set to 1 for MQA or any divisor of heads for GQA."
+        ),
+    )
+    dropout: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Dropout probability applied to the output projection, during training only.",
+    )
+    qk_norm: bool = Field(
+        default=False,
+        description=(
+            "Apply RMSNorm to queries and keys (in head space) before the dot product."
+            " Stabilises logit scale at large depth or large dim_head."
+        ),
+    )
+
+    @model_validator(mode="after")
+    def _validate(self) -> "CrossAttentionConfig":
+        _validate_attn_dims(self.dim, self.heads, self.dim_head, self.kv_heads)
+        return self
+
+    @property
+    def effective_kv_heads(self) -> int:
+        return self.kv_heads if self.kv_heads is not None else self.heads
+
+    @property
+    def groups(self) -> int:
+        return self.heads // self.effective_kv_heads

stackformers/attention/cross_attn.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import torch.nn as nn
+from einops import rearrange, repeat
+from torch import Tensor
+
+from stackformers.attention.config import CrossAttentionConfig
+from stackformers.attention.ops import packed_attn_or_fallback, padded_sdpa
+from stackformers.positional.protocols import PosEncoding
+from stackformers.sequence import PackedInput, PackedSequence, PaddedInput, SequenceInput
+
+
+class CrossAttention(nn.Module):
+    """Multi-head cross-attention: queries from x, keys/values from context.
+
+    Always global (no windowing). Pass PaddedInput for inference or export, PackedInput for
+    training. Falls back to padded SDPA when varlen_attn is unavailable (CPU, float32,
+    torch.export).
+    """
+
+    def __init__(self, config: CrossAttentionConfig, pos_encoding: PosEncoding) -> 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.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, x_input: PaddedInput, ctx_input: PaddedInput) -> Tensor:
+        cfg = self.config
+        h, kv_h, groups = cfg.heads, cfg.effective_kv_heads, cfg.groups
+        x, context = x_input.x, ctx_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(context), "b s (h d) -> b h s d", h=kv_h))
+        v = rearrange(self.to_v(context), "b s (h d) -> b h s d", h=kv_h)
+        if groups > 1:
+            k = repeat(k, "b h s d -> b (h g) s d", g=groups)
+            v = repeat(v, "b h s d -> b (h g) s d", g=groups)
+        q, k = self.pos_encoding.forward_padded(
+            q, k, x_input.abs_positions, ctx_input.abs_positions
+        )
+        out = padded_sdpa(q, k, v, ctx_input.mask, causal=False, window_size=None, bias=None)
+        out = self.dropout(self.to_out(rearrange(out, "b h n d -> b n (h d)")))
+        return out * x_input.mask.unsqueeze(-1)
+
+    def _forward_packed(self, x_input: PackedInput, ctx_input: PackedInput) -> Tensor:
+        cfg = self.config
+        h, kv_h, groups = cfg.heads, cfg.effective_kv_heads, cfg.groups
+        x, context = x_input.x, ctx_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(context), "nt (h d) -> nt h d", h=kv_h))
+        v = rearrange(self.to_v(context), "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, x_input.abs_positions, ctx_input.abs_positions
+        )
+        x_seq = PackedSequence(cu_seqlens=x_input.cu_seqlens, max_seqlen=x_input.max_seqlen)
+        ctx_seq = PackedSequence(cu_seqlens=ctx_input.cu_seqlens, max_seqlen=ctx_input.max_seqlen)
+        out = packed_attn_or_fallback(
+            q, k, v, x_seq, ctx_seq, causal=False, window_size=None, bias=None
+        )
+        return self.dropout(self.to_out(rearrange(out, "nt h d -> nt (h d)")))
+
+    def forward(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor:
+        match x_input:
+            case PaddedInput():
+                return self._forward_padded(x_input, ctx_input)  # type: ignore[arg-type]
+            case PackedInput():
+                return self._forward_packed(x_input, ctx_input)  # type: ignore[arg-type]

stackformers/attention/ops.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import warnings
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn.attention.varlen import varlen_attn as _varlen_attn
+
+from stackformers.sequence import PackedSequence
+
+
+def varlen_supported(q: Tensor) -> bool:
+    return q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)
+
+
+def padding_mask(mask: Tensor, dtype: torch.dtype) -> Tensor:
+    bias = torch.zeros(mask.shape, dtype=dtype, device=mask.device)
+    bias.masked_fill_(~mask, torch.finfo(dtype).min)
+    return bias.view(mask.shape[0], 1, 1, mask.shape[1])
+
+
+def window_mask(n: int, s: int, window_size: int, causal: bool, device: torch.device) -> Tensor:
+    """Additive sliding-window mask (0 = attend, -inf = ignore): (1, 1, n, s)."""
+    q_pos = torch.arange(n, device=device).unsqueeze(1)
+    k_pos = torch.arange(s, device=device).unsqueeze(0)
+    if causal:
+        allowed = (k_pos <= q_pos) & (k_pos >= q_pos - window_size)
+    else:
+        half = window_size // 2
+        allowed = (k_pos >= q_pos - half) & (k_pos <= q_pos + half)
+    mask = torch.zeros(1, 1, n, s, dtype=torch.float, device=device)
+    mask.masked_fill_(~allowed.unsqueeze(0).unsqueeze(0), float("-inf"))
+    return mask
+
+
+def varlen_window(causal: bool, window_size: int | None) -> tuple[int, int]:
+    if window_size is None:
+        return (-1, 0) if causal else (-1, -1)
+    return (window_size, 0) if causal else (window_size // 2, window_size // 2)
+
+
+def padded_sdpa(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    causal: bool,
+    window_size: int | None,
+    bias: Tensor | None,
+) -> Tensor:
+    """SDPA for padded inputs, applying padding mask + optional window mask + attention bias."""
+    n, s = q.shape[-2], k.shape[-2]
+    attn_mask = padding_mask(mask, q.dtype)
+    if bias is not None:
+        attn_mask = attn_mask + bias
+    if window_size is None:
+        if causal:
+            attn_mask = attn_mask + window_mask(n, s, s, causal=True, device=q.device)
+        # is_causal=False: causal constraint already encoded in attn_mask above
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, is_causal=False)
+    else:
+        win_mask = window_mask(n, s, window_size, causal, q.device)
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask + win_mask)
+
+
+def _cu_to_indices(cu: Tensor, b: int) -> tuple[Tensor, Tensor]:
+    """Return ``(batch_idx, pos_idx)``, each shape ``(nt,)``, from ``cu_seqlens``.
+
+    Both tensors are built without Python-side iteration over data-dependent lengths, so
+    this function is compatible with ``torch.export`` and ``torch.compile``.
+
+    The identity used for ``pos_idx`` is::
+
+        pos_idx[i] = i_global - cu[batch_idx[i]]
+
+    i.e. the within-document offset of the i-th flat token.
+    """
+    lengths = cu[1:] - cu[:-1]  # (b,) — never converted to a Python list
+    batch_idx = torch.repeat_interleave(
+        torch.arange(b, device=cu.device, dtype=torch.long), lengths
+    )  # (nt,)
+    nt = batch_idx.shape[0]
+    pos_idx = torch.arange(nt, device=cu.device, dtype=torch.long) - cu[batch_idx]
+    return batch_idx, pos_idx
+
+
+def _packed_heads_to_padded(x: Tensor, cu: Tensor, b: int, n: int) -> tuple[Tensor, Tensor]:
+    """Scatter (nt, h, d) → (b, h, n, d). Returns padded tensor and (b, n) bool mask."""
+    batch_idx, pos_idx = _cu_to_indices(cu, b)
+    h, d = x.shape[1], x.shape[2]
+    out = x.new_zeros(b, h, n, d)
+    out[batch_idx, :, pos_idx] = x
+    mask = torch.zeros(b, n, dtype=torch.bool, device=cu.device)
+    mask[batch_idx, pos_idx] = True
+    return out, mask
+
+
+def _padded_heads_to_packed(x: Tensor, mask: Tensor) -> Tensor:
+    """Gather (b, h, n, d) → (nt, h, d) using bool mask (b, n)."""
+    return x.permute(0, 2, 1, 3)[mask]
+
+
+def packed_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    q_seq: PackedSequence,
+    k_seq: PackedSequence,
+    causal: bool,
+    window_size: int | None,
+) -> Tensor:
+    win = varlen_window(causal, window_size)
+    result = _varlen_attn(
+        query=q,
+        key=k,
+        value=v,
+        cu_seq_q=q_seq.cu_seqlens.to(torch.int32),
+        cu_seq_k=k_seq.cu_seqlens.to(torch.int32),
+        max_q=q_seq.max_seqlen,
+        max_k=k_seq.max_seqlen,
+        window_size=win,
+    )
+    assert isinstance(result, Tensor)
+    return result
+
+
+def packed_attn_or_fallback(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    q_seq: PackedSequence,
+    k_seq: PackedSequence,
+    causal: bool,
+    window_size: int | None,
+    bias: Tensor | None,
+) -> Tensor:
+    """Run varlen attention; fall back to padded SDPA when unavailable.
+
+    Inputs and output are all packed: q/k/v are (nt, h, d), result is (nt_q, h, d).
+    bias=None → no attention bias; varlen_attn is tried first.
+    bias=Tensor → attention bias present; varlen_attn is skipped (it has no bias slot) and a
+    warning is emitted. Fallback also activates on CPU, non-float16/bfloat16 dtypes, or GPUs
+    that do not support varlen_attn (e.g. compute capability < 8.0).
+    """
+    if bias is not None and varlen_supported(q):
+        warnings.warn(
+            "varlen_attn does not support attention bias; falling back to padded SDPA. "
+            "Performance may be lower.",
+            stacklevel=2,
+        )
+    elif varlen_supported(q):
+        try:
+            return packed_attn(q, k, v, q_seq, k_seq, causal, window_size)
+        except RuntimeError as exc:
+            warnings.warn(
+                f"varlen_attn is not supported on this device ({exc}); "
+                "falling back to padded SDPA. Performance may be lower.",
+                stacklevel=2,
+            )
+
+    b = int(q_seq.cu_seqlens.shape[0]) - 1
+    q_pad, q_mask = _packed_heads_to_padded(q, q_seq.cu_seqlens, b, q_seq.max_seqlen)
+    k_pad, k_mask = _packed_heads_to_padded(k, k_seq.cu_seqlens, b, k_seq.max_seqlen)
+    v_pad, _ = _packed_heads_to_padded(v, k_seq.cu_seqlens, b, k_seq.max_seqlen)
+    out_pad = padded_sdpa(q_pad, k_pad, v_pad, k_mask, causal, window_size, bias)
+    return _padded_heads_to_packed(out_pad, q_mask)

stackformers/attention/protocols.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.sequence import PaddedInput, SequenceInput
+
+
+@runtime_checkable
+class AttnBias(Protocol):
+    """Additive bias injected into attention logits before softmax.
+
+    Receives the padded sequence input so implementations can use abs_positions.
+    Returns a broadcastable bias (1|b, 1|h, n, s), or None if the bias is a no-op
+    (which allows the packed path to use varlen_attn instead of SDPA).
+    Null implementation: NoAttnBias.
+    """
+
+    def __call__(
+        self,
+        input: PaddedInput,
+    ) -> Float[Tensor, "b h n s"] | None: ...
+
+
+@runtime_checkable
+class SelfAttn(Protocol):
+    """Self-attention over a sequence: maps input → x.
+
+    Implementation: SelfAttention.
+    """
+
+    def __call__(self, input: SequenceInput) -> Tensor: ...
+
+
+@runtime_checkable
+class CrossAttn(Protocol):
+    """Cross-attention from x to context: maps (x_input, ctx_input) → x.
+
+    Implementation: CrossAttention.
+    """
+
+    def __call__(self, x_input: SequenceInput, ctx_input: SequenceInput) -> Tensor: ...