stackformers 3.9.1__tar.gz → 4.2.0__tar.gz

stackformers-4.2.0/CHANGELOG.md

@@ -0,0 +1,320 @@
+# Changelog
+
+All notable changes to this project are documented here.
+
+The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project
+adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
+MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
+
+## [4.2.0] — 2026-07-14
+
+### Added
+
+- **`RotaryEmbeddingND`** (`positional/rope_nd.py`) — rotary position encoding over `c` spatial
+  dimensions, splitting `dim_head` into `c` per-axis blocks. `c=1` and `c=2` reduce to the
+  existing 1-D and 2-D encodings; 3-D and beyond now work with the same class.
+- `RoPENDConfig`, parameterised by **`r_min` / `r_max` instead of `base`**.
+
+### Notes
+
+- **`base` is wrong for continuous coordinates**, and `RoPENDConfig` therefore does not have
+  it. RoPE's fastest band is pinned at `ω = 1` — a wavelength of exactly 2π *coordinate units* —
+  whatever the base; `base` can only stretch the slow end. So the ladder lands correctly only
+  when tokens happen to sit one unit apart. That is true of text and of a patch grid, and of
+  nothing else: `base=10000` on a 14×14 patch grid leaves 14 of its 16 bands frozen, and
+  coordinates normalised to `[0, 1]` leave every band longer than the whole domain.
+
+  The ladder is built from the data instead — `ω_hi = π / r_min` (Nyquist on the finest
+  separation the model must resolve) down to `ω_lo = 2π / (headroom · r_max)` (the longest
+  wavelength spans the domain). It then depends only on the *dynamic range* `r_max / r_min`, so
+  metres, pixels and millimetres give the identical encoding — the property `base` never had,
+  and the one that makes the setting measurable from the data rather than tuned.
+
+- `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged. Use them for text and grids.
+
+- **Centre your coordinates.** Not for invariance — RoPE is translation-invariant exactly, since
+  the rotations cancel into `ω·(pᵢ − pⱼ)`. For float32: the angle is `ω·p`, and coordinates far
+  from the origin push it into the hundreds of radians and spend the mantissa before the cosine
+  is taken. A 1e5 offset costs three orders of magnitude of accuracy.
+
+- **It encodes direction, not just distance**, so it is not rotation invariant — by design, as
+  the relative offset is strictly more information than the distance. Where the global frame is
+  arbitrary, train with rotation augmentation. Exact rotation invariance in the architecture
+  costs either an O(n²) bias (`RelativeDistanceBias`, above) or a canonicalisation of the input
+  frame, which is inherently discontinuous: a near-isotropic point set flips frames under a 1%
+  perturbation, measured on 1–2% of samples.
+
+## [4.1.0] — 2026-07-14
+
+### Added
+
+- **`RelativeDistanceBias`** (`attention/distance_bias.py`) — an additive attention bias read
+  off the Euclidean distance between node positions: pairwise distance → Gaussian radial
+  shells → a learned per-head profile. Only `‖pᵢ − pⱼ‖` enters the logit, so attention is
+  invariant to any global translation *or rotation* of the node set. For 2-D/N-D node sets
+  whose coordinate frame is arbitrary.
+- `DistanceBiasConfig`, `NoAttnBiasConfig`, and the `AttnBiasConfig` discriminated union;
+  `attention/factory.py` with `build_attn_bias`; a `node_encoder_config` preset pairing the
+  bias with `NoPosEncoding`.
+
+### Notes
+
+- **A rotary encoding cannot express this.** RoPE rotates by `ω·p`, which is linear in position
+  by construction — and that linearity is exactly what makes the query and key rotations cancel
+  into a relative offset. Distance is not linear in position, so it can only enter as a bias.
+  Concretely, RoPE-2D scores two neighbours *the same distance away* 2.7× differently depending
+  on whether the offset is axis-aligned or diagonal, and rotating a node set by 30° moves its
+  attention logits by 83%.
+- **Opt-in, and it has a ceiling.** `NoAttnBias` remains the default and returns `None`, which
+  keeps `varlen_attn` available, so the bias costs nothing when unused. When used it costs a
+  great deal: a bias has no varlen slot, so it forces padded SDPA and an O(n²) tensor, and the
+  `(b, n, s, num_rbf)` shell intermediate — retained for backward, and `num_rbf/h` times larger
+  than the bias itself — dominates activation memory. Measured OOM past roughly 2 000 nodes on
+  16 GiB with 4 layers. Where that does not fit, `RotaryEmbeddingND` plus rotation augmentation
+  is the cheaper trade.
+
+## [4.0.1] — 2026-07-14
+
+### Fixed
+
+- **Causal and windowed attention crashed on CUDA in half precision.** `window_mask` built its
+  mask at `torch.float` while `padding_mask` used the query's dtype; adding them promoted the
+  sum to float32, and CUDA's SDPA rejects a mask whose dtype differs from the query
+  (`RuntimeError: invalid dtype for bias`). Every causal or sliding-window call in
+  `float16`/`bfloat16` was affected — that is, the configuration a model actually trains in.
+  `window_mask` now takes an explicit `dtype` and `padded_sdpa` passes `q.dtype`.
+
+  It went unnoticed because CPU's SDPA math backend tolerates the mismatch, and the suite had
+  only ever run on CPU — half precision exists solely in the CUDA half of the device matrix.
+  Running the untouched tree on a GPU fails 30 tests that pass locally.
+
+## [4.0.0] — 2026-07-14
+
+### Removed
+
+- **Breaking: the 2-D vision modules.** The `spatial` package (`WindowAttention2D`,
+  `SpatialReductionAttention`, `ConvKVReduction`/`NoKVReduction`, `PatchMerging`,
+  `SpatialInput`, `SpatialTransformerLayer`) and the `PyramidVisionBackbone` preset with
+  `pyramid_vision_config` are gone. The library returns to sequence models only.
+- `RotaryEmbedding2D` and `RoPE2DConfig` are **kept** — they are independent of `spatial` and
+  encode any two-coordinate position (grid row/col, or point x/y).
+
+### Notes
+
+- A coordinate-frame bug in `ConvKVReduction` (pooled keys were positioned in the reduced
+  grid's index space while queries used full-grid coordinates, so a relative encoding saw an
+  offset that grew with absolute position) was found but not fixed, as the module is removed
+  here. Anyone reviving this code from history must fix it first: pooled tokens belong at their
+  block centre, `i·r + (r − 1)/2`, in the original grid's frame.
+
+## [3.9.2] — 2026-07-14
+
+### Fixed
+
+- **`RotaryEmbedding2D` was not computing a rotation.** The frequency vector was laid out as
+  `[row | row | col | col]`, but `_rotate_half` pairs channel `i` with channel `i + dh/2`, so a
+  row angle was paired against a column angle. The resulting 2×2 map had
+  `det = cos((row − col) · ωᵢ)` and singular values `√(1 ± |sin D|)` — a squeeze, not a rotation.
+  It was therefore non-orthogonal, annihilated high-frequency channels wherever the determinant
+  reached zero, and produced attention scores that leaked absolute position instead of depending
+  only on relative offset. Frequencies are now built as `[row | col | row | col]`.
+- `RotaryEmbedding2D` now casts positions and `inv_freq` to float32 for the outer product,
+  matching `RotaryEmbedding1D`. A `.half()`-ed module previously lost position precision.
+
+### Changed
+
+- **Numerics of `RotaryEmbedding2D` change.** Any checkpoint trained against the previous
+  encoding is tuned to the broken position code and requires retraining or finetuning.
+
+### Added
+
+- Math tests for `RotaryEmbedding2D` on a genuine row-major grid: a definition-first reference
+  oracle, orthogonality (norms and inner products), relative-offset invariance, per-axis
+  sensitivity, distinct-offset separation, and packed/padded agreement. The previous test used a
+  `row == col` position grid — the main diagonal, the one locus where the row and column angles
+  coincide and the broken layout is indistinguishable from a correct one.
+
+## [3.9.1] — 2026-06-30
+
+### Fixed
+
+- Make the experimental `varlen_attn` import and call robust across PyTorch versions.
+
+## [3.9.0] — 2026-06-17
+
+### Added
+
+- 2-D grid attention (`spatial`) and the pyramid vision backbone preset.
+
+## [3.8.2] — 2026-05-29
+
+### Fixed
+
+- Force absolute positions to float32 in `sequence`.
+
+## [3.8.1] — 2026-05-29
+
+### Fixed
+
+- Keep batch size as a `SymInt` in the packed attention fallback, preserving `torch.export`
+  compatibility.
+
+## [3.8.0] — 2026-05-27
+
+### Fixed
+
+- Replace `.tolist()` loops in `sequence` and `ops` with tensor operations, making them
+  `torch.export`-compatible.
+
+### Added
+
+- PyPI publish metadata in `pyproject.toml`.
+
+## [3.7.0] — 2026-05-21
+
+### Changed
+
+- Introduce abstract generic base classes for all presets.
+
+## [3.6.0] — 2026-05-21
+
+### Added
+
+- `AttnBias` protocol and the `NoAttnBias` null object on `SelfAttention`.
+
+## [3.5.2] — 2026-05-21
+
+### Fixed
+
+- Warn and fall back to SDPA when `varlen_attn` raises.
+
+## [3.5.1] — 2026-05-18
+
+### Changed
+
+- Move attention dropout to the output projection.
+
+## [3.5.0] — 2026-05-18
+
+### Added
+
+- Packed→padded SDPA fallback; extract `ops.py`.
+
+## [3.4.2] — 2026-05-18
+
+### Added
+
+- `padded_to_packed` and `packed_to_padded` helpers.
+
+## [3.4.1] — 2026-05-18
+
+### Removed
+
+- Packed-attention CPU fallback; warn on dropout instead.
+
+## [3.4.0] — 2026-05-17
+
+### Added
+
+- ReLU² feed-forward network.
+
+## [3.3.0] — 2026-05-17
+
+### Added
+
+- QK-norm in attention; all config fields documented.
+
+## [3.2.0] — 2026-05-17
+
+### Changed
+
+- Eliminate layout-specific duplicates in positional encoding implementations.
+
+## [3.1.0] — 2026-05-17
+
+### Changed
+
+- Replace `forward()` on positional encodings with `forward_padded` / `forward_packed`, so
+  callers pick the path for the layout they already know.
+
+## [3.0.1] — 2026-05-17
+
+### Fixed
+
+- Use the GELU tanh approximation in GEGLU.
+
+## [3.0.0] — 2026-05-16
+
+### Changed
+
+- **Breaking:** remove the kernel abstraction; padded and packed dispatch are unified.
+
+## [2.0.2] — 2026-05-16
+
+### Removed
+
+- The unused local-attention package.
+
+## [2.0.1] — 2026-05-16
+
+### Changed
+
+- Store feed-forward activations as `nn.Module` attributes.
+
+## [2.0.0] — 2026-05-16
+
+### Removed
+
+- **Breaking:** the `attn_bias` / ALiBi system.
+
+## [1.1.0] — 2026-05-16
+
+### Added
+
+- `LearnedPosEncoding` with absolute position embeddings.
+
+## [1.0.1] — 2026-05-16
+
+### Added
+
+- `LayerNorm` wrapper.
+
+## [1.0.0] — 2026-05-16
+
+### Added
+
+- GEGLU feed-forward variant.
+
+### Changed
+
+- **Breaking:** feed-forward config becomes a discriminated union.
+
+## [0.5.0] — 2026-05-16
+
+### Added
+
+- Unfold mode for the windowed SDPA kernel.
+
+## [0.4.0] — 2026-05-15
+
+### Added
+
+- `RotaryEmbedding2D` and the packed cross-attender.
+
+## [0.3.0] — 2026-05-14
+
+### Changed
+
+- Establish the semantic versioning convention.
+
+## [0.2.0] — 2026-05-13
+
+### Changed
+
+- Remove the `v1` namespace; add the `presets` package.
+
+## [0.1.0] — 2026-05-13
+
+### Added
+
+- Initial library scaffold.

{stackformers-3.9.1 → stackformers-4.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 3.9.1
+Version: 4.2.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

{stackformers-3.9.1 → stackformers-4.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "3.9.1"
+version = "4.2.0"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

stackformers-4.2.0/run-remote-tests.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# Run the test suite on a remote CUDA box.
+#
+# flex_attention's backward is CUDA-only, so the tests that cover it are skipped on this
+# Mac and only ever execute here.
+#
+#   ./run-remote-tests.sh                          # whole suite
+#   ./run-remote-tests.sh tests/attention -k flex  # anything after the script is passed to pytest
+#   REMOTE=other-box ./run-remote-tests.sh
+set -euo pipefail
+
+REMOTE="${REMOTE:-srukr-redeyed-pc}"
+REMOTE_DIR="${REMOTE_DIR:-/tmp/sf-cuda/repo}"
+UV="${UV:-\$HOME/.local/bin/uv}" # resolved on the remote, not here
+LOCAL_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+echo "==> syncing $LOCAL_DIR -> $REMOTE:$REMOTE_DIR"
+rsync -az --delete \
+    --exclude '.venv' --exclude '.git' --exclude 'x-transformers' \
+    --exclude '__pycache__' --exclude '.pytest_cache' --exclude '.ruff_cache' \
+    "$LOCAL_DIR/" "$REMOTE:$REMOTE_DIR/"
+
+echo "==> uv sync + pytest on $REMOTE"
+ssh "$REMOTE" "cd '$REMOTE_DIR' && $UV sync --quiet && $UV run pytest ${*:-tests/} -q"

{stackformers-3.9.1 → stackformers-4.2.0}/stackformers/attention/README.md

@@ -18,6 +18,14 @@ Multi-head self-attention and cross-attention.
 
 **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.
 
+**Geometry a rotary encoding cannot carry.** `RelativeDistanceBias` adds a learned function of the Euclidean distance between node positions to the logits, making attention invariant to any global translation or rotation of the input. It lives here rather than in `positional/` because it *cannot* live there: RoPE rotates by `ω·p`, which is linear in position by construction — and that linearity is exactly what makes the query and key rotations cancel into a relative offset. Distance is not linear in position, so no rotary encoding, axial or otherwise, can express it; a bias is the only route. Pair it with `NoPosEncoding` — a rotary encoding alongside it would reintroduce the preferred frame it exists to remove. Use it when node coordinates have no meaningful axes; when up and right *do* mean something, RoPE-2D keeps direction and stays varlen-compatible.
+
+**The bias costs the varlen path, by construction.** `varlen_attn` has no bias slot, so any non-`None` `AttnBias` forces the padded SDPA path and materialises a `(b, h, n, s)` tensor. The `(b, n, s, num_rbf)` shell intermediate is retained for backward and is `num_rbf/h` times larger still — it, not the bias, dominates activation memory. `NoAttnBias` returns `None` precisely so the common case keeps the kernel. `TransformerEncoder` shares one bias *module* across layers, but each layer still *calls* it, so the tensor is recomputed and retained per layer; at large node counts the levers are a smaller `num_rbf`, checkpointing the bias, or hoisting the call out of the layer loop.
+
+**Why not `flex_attention`.** It fuses the bias into the kernel and would make it O(n) instead of O(n²) — attempted, and parked on the `flex-attention` branch. It needs `dim_head >= 16`; its `score_mod` must lower as a *pointwise* subgraph, which forbids both a `vector_norm` over coordinates and a sum over shells; a captured tensor requiring grad may be indexed exactly once, which forbids unrolling that sum instead; called eagerly it is unfused and catastrophic (measured 64 s/step at n=2048, OOM at n=4096, against ~40 ms for SDPA), so it is viable only inside `torch.compile`; and under `torch.compile` with *dynamic* shapes it bails back to SDPA anyway. With node counts that vary per batch, the fused kernel would rarely be the one that runs.
+
 ## 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.
+
+To add a new attention bias, write an `nn.Module` satisfying `AttnBias` (take `PaddedInput`, return a broadcastable `(b, h, n, s)` tensor or `None`), give it a config with a `kind` discriminator, join it to the `AttnBiasConfig` union, and add one `match` arm to `build_attn_bias`.

{stackformers-3.9.1 → stackformers-4.2.0}/stackformers/attention/__init__.py

@@ -1,16 +1,29 @@
 from stackformers.attention.bias import NoAttnBias
-from stackformers.attention.config import CrossAttentionConfig, SelfAttentionConfig
+from stackformers.attention.config import (
+    AttnBiasConfig,
+    CrossAttentionConfig,
+    DistanceBiasConfig,
+    NoAttnBiasConfig,
+    SelfAttentionConfig,
+)
 from stackformers.attention.cross_attn import CrossAttention
+from stackformers.attention.distance_bias import RelativeDistanceBias
+from stackformers.attention.factory import build_attn_bias
 from stackformers.attention.protocols import AttnBias, CrossAttn, SelfAttn
 from stackformers.attention.self_attn import SelfAttention
 
 __all__ = [
     "SelfAttentionConfig",
     "CrossAttentionConfig",
+    "AttnBiasConfig",
+    "NoAttnBiasConfig",
+    "DistanceBiasConfig",
     "SelfAttn",
     "CrossAttn",
     "AttnBias",
     "SelfAttention",
     "CrossAttention",
     "NoAttnBias",
+    "RelativeDistanceBias",
+    "build_attn_bias",
 ]

{stackformers-3.9.1 → stackformers-4.2.0}/stackformers/attention/config.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field, model_validator
 
@@ -96,6 +97,48 @@ class SelfAttentionConfig(BaseModel):
         return self.heads // self.effective_kv_heads
 
 
+class NoAttnBiasConfig(BaseModel):
+    """No attention bias — leaves varlen_attn available on the packed path."""
+
+    kind: Literal["none"] = "none"
+
+
+class DistanceBiasConfig(BaseModel):
+    """Config for RelativeDistanceBias.
+
+    Attention logits gain a learned function of the Euclidean distance between node
+    positions, making them invariant to global translation and rotation of the node set.
+    Use with NoPosEncodingConfig: a rotary encoding would reintroduce a preferred frame.
+    """
+
+    kind: Literal["distance"] = "distance"
+    heads: int = Field(gt=0, description="Number of query heads. Must match the attention config.")
+    r_max: float = Field(
+        gt=0.0,
+        description=(
+            "Largest pairwise node distance to resolve. Radial-basis shells are spread over"
+            " [0, r_max]; pairs further apart fall in the tail of the outermost shell and become"
+            " indistinguishable. Measure it from the data — a high percentile of the pairwise"
+            " distance distribution, not the maximum, which is an outlier."
+        ),
+    )
+    num_rbf: int = Field(
+        default=32,
+        gt=1,
+        description=(
+            "Number of Gaussian radial-basis shells. Distance resolution is r_max / (num_rbf - 1);"
+            " raising it sharpens the profile a head can learn but costs a (b, n, s, num_rbf)"
+            " intermediate, which dominates activation memory at large node counts."
+        ),
+    )
+
+
+AttnBiasConfig = Annotated[
+    NoAttnBiasConfig | DistanceBiasConfig,
+    Field(discriminator="kind"),
+]
+
+
 class CrossAttentionConfig(BaseModel):
     dim: int = Field(
         gt=0,

stackformers-4.2.0/stackformers/attention/distance_bias.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.attention.config import DistanceBiasConfig
+from stackformers.sequence import PaddedInput
+
+
+class RelativeDistanceBias(nn.Module):
+    """Additive attention bias read off the Euclidean distance between node positions.
+
+    Only ‖p_i − p_j‖ enters the logit, so attention is invariant to any global translation
+    or rotation of the node set: the encoding privileges no axis and no origin. Each head
+    mixes a shared bank of Gaussian radial-basis shells into its own profile over distance,
+    so some heads can go local and others long-range.
+
+    A rotary encoding cannot express this. RoPE rotates by ω·p, which is linear in position
+    by construction — that linearity is exactly what makes the query and key rotations cancel
+    into a relative offset. Distance is not linear in position, so it can only enter as a bias.
+
+    Costs one (b, h, n, s) tensor, which rules out varlen_attn (it has no bias slot) and forces
+    the padded SDPA path. Affordable for node counts in the low thousands; beyond that the
+    (b, n, s, num_rbf) shell intermediate, not the bias, is what dominates activation memory.
+    """
+
+    def __init__(self, config: DistanceBiasConfig) -> None:
+        super().__init__()
+        centres = torch.linspace(0.0, config.r_max, config.num_rbf)
+        self.register_buffer("centres", centres, persistent=False)
+        self.width = config.r_max / (config.num_rbf - 1)  # neighbouring shells cross at 1/e
+        self.to_bias = nn.Linear(config.num_rbf, config.heads, bias=False)
+        nn.init.zeros_(self.to_bias.weight)  # start at content-only attention, learn the profile
+
+    def _shell_weights(self, dist: Float[Tensor, "b n s"]) -> Float[Tensor, "b n s k"]:
+        centres: Tensor = self.centres  # type: ignore[assignment]
+        return torch.exp(-(((dist.unsqueeze(-1) - centres) / self.width) ** 2))
+
+    def forward(self, input: PaddedInput) -> Float[Tensor, "b h n s"]:
+        """float32 distances keep half-precision coordinates from collapsing nearby nodes."""
+        pos = input.abs_positions.float()
+        delta = pos.unsqueeze(2) - pos.unsqueeze(1)  # b n s c
+        dist = torch.linalg.vector_norm(delta, dim=-1)  # b n s
+        shells = self._shell_weights(dist).to(self.to_bias.weight.dtype)
+        bias = rearrange(self.to_bias(shells), "b n s h -> b h n s")
+        return bias.to(input.x.dtype)

stackformers-4.2.0/stackformers/attention/factory.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from stackformers.attention.bias import NoAttnBias
+from stackformers.attention.config import (
+    AttnBiasConfig,
+    DistanceBiasConfig,
+    NoAttnBiasConfig,
+)
+from stackformers.attention.distance_bias import RelativeDistanceBias
+from stackformers.attention.protocols import AttnBias
+
+
+def build_attn_bias(config: AttnBiasConfig) -> AttnBias:
+    match config:
+        case NoAttnBiasConfig():
+            return NoAttnBias()
+        case DistanceBiasConfig():
+            return RelativeDistanceBias(config)
+        case _:
+            raise AssertionError(f"Unhandled attn bias config: {type(config)}")

{stackformers-3.9.1 → stackformers-4.2.0}/stackformers/attention/ops.py

@@ -14,8 +14,16 @@ def padding_mask(mask: Tensor, dtype: torch.dtype) -> Tensor:
     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)."""
+def window_mask(
+    n: int, s: int, window_size: int, causal: bool, device: torch.device, dtype: torch.dtype
+) -> Tensor:
+    """Additive sliding-window mask (0 = attend, -inf = ignore): (1, 1, n, s).
+
+    ``dtype`` must be the query's. CUDA's SDPA rejects a mask whose dtype differs from the
+    query ("invalid dtype for bias"), and a float32 mask added to a float16 padding mask
+    promotes the sum to float32 — which broke every causal and windowed call in half
+    precision on GPU, while passing on CPU, whose math backend tolerates the mismatch.
+    """
     q_pos = torch.arange(n, device=device).unsqueeze(1)
     k_pos = torch.arange(s, device=device).unsqueeze(0)
     if causal:
@@ -23,7 +31,7 @@ def window_mask(n: int, s: int, window_size: int, causal: bool, device: torch.de
     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 = torch.zeros(1, 1, n, s, dtype=dtype, device=device)
     mask.masked_fill_(~allowed.unsqueeze(0).unsqueeze(0), float("-inf"))
     return mask
 
@@ -44,11 +52,11 @@ def padded_sdpa(
         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)
+            attn_mask = attn_mask + window_mask(n, s, s, True, q.device, q.dtype)
         # 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)
+        win_mask = window_mask(n, s, window_size, causal, q.device, q.dtype)
         return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask + win_mask)
 
 

stackformers-4.2.0/stackformers/positional/README.md

@@ -0,0 +1,28 @@
+# 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` takes two coordinates per token (a grid's row/col, or a point's x/y — the two axes are treated symmetrically, so the order is yours to pick as long as it is consistent) and splits the head dimension between them. `NoPosEncoding` is a null object for cross-attention paths that need no positional information.
+
+## `base` is for grids. `RotaryEmbeddingND` is for continuous coordinates.
+
+RoPE's shortest wavelength is fixed at 2π coordinate units, and `base` only stretches the long end of the band range. So `base` lands the ladder correctly only when tokens sit *one unit apart* — true for text, true for a patch grid, meaningless for scattered continuous coordinates. Positions normalised to `[0, 1]` leave every band longer than the whole domain, so the encoding barely rotates and carries almost no signal; and `base=10000` on a 14×14 patch grid leaves 14 of 16 bands frozen.
+
+`RotaryEmbeddingND` (`RoPENDConfig`) drops `base` entirely and builds the ladder from the two numbers that actually mean something, in any number of dimensions:
+
+- `r_min` — the finest separation the model must resolve. Sets the shortest wavelength to `2·r_min`, the Nyquist limit. Measure it as a low percentile of the nearest-neighbour distance.
+- `r_max` — the domain diameter. Sets the longest wavelength (times `headroom`, so the slowest band stays monotone rather than wrapping). Measure it as a high percentile of the pairwise distance distribution.
+
+The resulting ladder depends only on the *dynamic range* `r_max / r_min`, so metres, pixels and millimetres all give the identical encoding — the property `base` never had. `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged; use them for text and grids.
+
+**Centre your coordinates.** Not for invariance — RoPE is translation-invariant exactly, since the query and key rotations cancel into `ω·(pᵢ − pⱼ)`. It is for float32: the angle is `ω·p`, and coordinates far from the origin push it into the hundreds of radians and spend the mantissa before the cosine is taken. A 1e5 offset costs three orders of magnitude of accuracy. This cannot be done inside the module — cross-attention must subtract the *same* constant from query and key positions, and the module sees them separately — so subtract the centroid at the input boundary.
+
+**RoPE encodes direction, not just distance.** It is not rotation invariant, by design: it gives the model the relative *offset*, which is strictly more information than the distance. Where the global frame is arbitrary, train with rotation augmentation. Buying exact rotation invariance in the architecture costs either an O(n²) attention bias (see `attention/RelativeDistanceBias`, which OOMs past a few thousand nodes) or a canonicalisation of the input frame, which is inherently discontinuous — a near-isotropic point set flips frames under a 1% perturbation.
+
+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`.