stackformers 3.9.0__tar.gz → 3.9.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 3.9.0
+Version: 3.9.1
 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.0 → stackformers-3.9.1}/pyproject.toml

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

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

@@ -8,9 +8,11 @@ Multi-head self-attention and cross-attention.
 
 **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.
+**Kernel ops in `ops.py`.** All low-level attention math lives in `attention/ops.py`: mask builders, `padded_sdpa`, the packed↔padded scatter/gather helpers, 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).
+**Experimental kernel behind one interface function.** `torch.nn.attention.varlen` is an experimental API — the symbol is absent in some torch builds and its signature is not stable. Everything that touches it lives in `varlen_backend.py` behind a single function, `try_varlen_attn(...) -> Tensor | None`: the guarded import, the eligibility/bias checks, and the call itself. The rest of the package depends only on that function and never on the volatile import path, so an absent or renamed symbol degrades to the fallback instead of breaking the whole package at import time.
+
+**Packed attention with automatic fallback.** `PackedInput` normally routes to `varlen_attn`, which requires CUDA and `float16`/`bfloat16`. `packed_attn_or_fallback` calls `try_varlen_attn`; a `None` result means take the padded path — scatter q/k/v to padded layout, run `F.scaled_dot_product_attention`, and gather the result back to packed (same weights, same output shape `(nt, h, d)`). The fallback runs **silently** for the expected CPU/float32/`torch.export` route, and **with a `UserWarning` explaining why** when the kernel was eligible but unusable: it failed to import, an attention bias is present (no bias slot), or the call raised a runtime/signature error (unsupported GPU or a changed experimental API). 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.
 

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

@@ -1,19 +1,13 @@
 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.attention.varlen_backend import try_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)
@@ -34,12 +28,6 @@ def window_mask(n: int, s: int, window_size: int, causal: bool, device: torch.de
     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,
@@ -101,30 +89,6 @@ def _padded_heads_to_packed(x: Tensor, mask: Tensor) -> Tensor:
     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,
@@ -135,29 +99,16 @@ def packed_attn_or_fallback(
     window_size: int | None,
     bias: Tensor | None,
 ) -> Tensor:
-    """Run varlen attention; fall back to padded SDPA when unavailable.
+    """Run varlen attention; fall back to padded SDPA when it is 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).
+    ``try_varlen_attn`` owns the experimental kernel and its warnings; ``None`` from it means
+    take the padded SDPA path — silently for the expected CPU/float32/export route, or after
+    a warning when the kernel was eligible but unavailable.
     """
-    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,
-            )
+    out = try_varlen_attn(q, k, v, q_seq, k_seq, causal, window_size, bias)
+    if out is not None:
+        return out
 
     # Keep b as SymInt-friendly: int(SymInt) forces specialization under torch.export.
     # Downstream consumers (_packed_heads_to_padded) accept SymInt for shape args.

stackformers-3.9.1/stackformers/attention/varlen_backend.py

@@ -0,0 +1,88 @@
+"""The experimental ``varlen_attn`` kernel behind one interface function.
+
+``torch.nn.attention.varlen`` is an experimental API: the symbol is absent in some torch
+builds, and its call signature is not yet stable. Everything that touches it lives here —
+the guarded import, the eligibility/bias checks, and the call itself — so the rest of the
+package depends only on :func:`try_varlen_attn` and never on the volatile import path.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+from torch import Tensor
+
+from stackformers.sequence import PackedSequence
+
+_FALLBACK_NOTE = " Falling back to padded SDPA; performance may be lower."
+
+
+def _load():
+    try:
+        from torch.nn.attention.varlen import varlen_attn
+    except ImportError as exc:
+        return None, f"varlen_attn could not be imported ({exc})"
+    return varlen_attn, None
+
+
+_varlen_attn, _import_error = _load()
+
+
+def _eligible(q: Tensor) -> bool:
+    """Whether this query's device/dtype is one ``varlen_attn`` would handle at all."""
+    return q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)
+
+
+def _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 _skip(reason: str) -> None:
+    warnings.warn(reason + _FALLBACK_NOTE, stacklevel=3)
+    return None
+
+
+def try_varlen_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    q_seq: PackedSequence,
+    k_seq: PackedSequence,
+    causal: bool,
+    window_size: int | None,
+    bias: Tensor | None,
+) -> Tensor | None:
+    """Attempt varlen attention for a packed batch; return ``None`` to use the SDPA fallback.
+
+    q/k/v are ``(nt, h, d)``; the result is the packed ``(nt_q, h, d)`` tensor. ``None`` means
+    the caller should run the padded SDPA path instead. When varlen *would* apply (CUDA +
+    float16/bfloat16) but cannot, a ``UserWarning`` explains why: an absent kernel, an
+    attention bias (varlen_attn has no bias slot), or a runtime/signature error from the call
+    (an unsupported GPU or a changed experimental API). The expected CPU/float32/export route
+    returns ``None`` silently — that is normal, not a failure.
+    """
+    if not _eligible(q):
+        return None
+    if _varlen_attn is None:
+        return _skip(f"varlen_attn is unavailable ({_import_error}).")
+    if bias is not None:
+        return _skip("varlen_attn does not support attention bias.")
+    try:
+        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=_window(causal, window_size),
+        )
+    except (RuntimeError, TypeError) as exc:
+        return _skip(f"varlen_attn call failed ({type(exc).__name__}: {exc}).")
+    if not isinstance(result, Tensor):
+        return _skip(f"varlen_attn returned {type(result).__name__}, expected Tensor.")
+    return result

stackformers-3.9.1/tests/attention/test_varlen_backend.py

@@ -0,0 +1,122 @@
+"""Tests for stackformers/attention/varlen_backend.py.
+
+Covers the fallback contract of ``try_varlen_attn``: it returns the kernel's tensor on the
+happy path and ``None`` (with a reason-bearing ``UserWarning`` when the kernel was eligible
+but unusable) on every degraded path — absent symbol, attention bias, a runtime/signature
+error from the call, or a non-Tensor return. The experimental kernel is faked via monkeypatch
+so the logic is exercised without CUDA or the real (possibly-absent) ``varlen_attn``.
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import pytest
+import torch
+
+from stackformers.attention import varlen_backend as vb
+from stackformers.sequence import PackedSequence
+
+NT, H, DH = 5, 2, 8
+
+Packed = tuple[torch.Tensor, torch.Tensor, torch.Tensor, PackedSequence]
+
+
+@pytest.fixture
+def packed() -> Packed:
+    """Two packed sequences (lengths 3 and 2) of (nt, h, d) q/k/v plus their PackedSequence."""
+    seq = PackedSequence(cu_seqlens=torch.tensor([0, 3, 5], dtype=torch.long), max_seqlen=3)
+    q = torch.randn(NT, H, DH)
+    k = torch.randn(NT, H, DH)
+    v = torch.randn(NT, H, DH)
+    return q, k, v, seq
+
+
+@pytest.fixture
+def force_eligible(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Pretend the query is CUDA + float16 so the eligibility guard does not short-circuit."""
+    monkeypatch.setattr(vb, "_eligible", lambda q: True)
+
+
+def _attempt(packed: Packed, bias: torch.Tensor | None = None) -> torch.Tensor | None:
+    q, k, v, seq = packed
+    return vb.try_varlen_attn(q, k, v, seq, seq, causal=False, window_size=None, bias=bias)
+
+
+def test_ineligible_falls_back_silently(packed: Packed, recwarn: pytest.WarningsRecorder) -> None:
+    """A plain CPU/float32 query is the expected fallback route — None, and no warning."""
+    out = _attempt(packed)
+    assert out is None
+    assert len(recwarn) == 0
+
+
+def test_unavailable_kernel_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """An eligible query with no imported kernel falls back and reports the import reason."""
+    monkeypatch.setattr(vb, "_varlen_attn", None)
+    monkeypatch.setattr(vb, "_import_error", "varlen_attn could not be imported (no module)")
+    with pytest.warns(UserWarning, match="unavailable"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_attention_bias_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """varlen_attn has no bias slot: an eligible query with a bias falls back with a warning."""
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: torch.zeros(NT, H, DH))
+    bias = torch.zeros(1, 1, NT, NT)
+    with pytest.warns(UserWarning, match="attention bias"):
+        out = _attempt(packed, bias=bias)
+    assert out is None
+
+
+def test_signature_error_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A changed experimental signature raises TypeError, which is caught and falls back."""
+
+    def renamed_kwarg(**kw: object) -> torch.Tensor:
+        raise TypeError("unexpected keyword argument 'cu_seq_q'")
+
+    monkeypatch.setattr(vb, "_varlen_attn", renamed_kwarg)
+    with pytest.warns(UserWarning, match="call failed"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_runtime_error_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """An unsupported-GPU RuntimeError from the call is caught and falls back."""
+
+    def unsupported(**kw: object) -> torch.Tensor:
+        raise RuntimeError("varlen_attn requires compute capability >= 8.0")
+
+    monkeypatch.setattr(vb, "_varlen_attn", unsupported)
+    with pytest.warns(UserWarning, match="call failed"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_non_tensor_return_warns(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A changed return contract (non-Tensor) falls back rather than propagating bad data."""
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: ("not", "a", "tensor"))
+    with pytest.warns(UserWarning, match="expected Tensor"):
+        out = _attempt(packed)
+    assert out is None
+
+
+def test_success_returns_kernel_tensor(
+    packed: Packed, force_eligible: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The happy path returns the kernel's tensor unchanged and emits no warning."""
+    expected = torch.randn(NT, H, DH)
+    monkeypatch.setattr(vb, "_varlen_attn", lambda **kw: expected)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        out = _attempt(packed)
+    assert out is expected

{stackformers-3.9.0 → stackformers-3.9.1}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "3.9.0"
+version = "3.9.1"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },