build(worker): port bidirectionalization seam to transformers 5.14.1

packages/a2d-worker-hf/pyproject.toml

@@ -6,10 +6,10 @@ dependencies = [
     "a2d-contracts",
     "pydantic>=2",
     "torch",
-    # Pinned: Decision 2 patches HF's eager attention seams directly. 4.51.3 is the
-    # first pin that ships Gemma 3 (`gemma3_text`); the GPT-2 `eager_attention_forward`
-    # signature and the Gemma 1 `_update_causal_mask` seam are unchanged from 4.48.3.
-    "transformers==4.51.3",
+    # Pinned: Decision 2 patches HF's eager attention seam directly, so a minor bump can
+    # silently move it. v5 replaced the per-model `_update_causal_mask` and GPT-2's
+    # `self.bias` with the `masking_utils` mask-interface; see `transform/attention.py`.
+    "transformers==5.14.1",
     "accelerate",
     "safetensors",
 ]

packages/a2d-worker-hf/src/a2d_core/transform/apply.py

@@ -13,7 +13,7 @@
 from typing import Any
 
 from a2d_core.transform.attention import AnnealState
-from a2d_core.transform.gqa_attention import _find_causal_mask_owner
+from a2d_core.transform.gqa_attention import is_rope_family
 from a2d_core.transform.handlers import TRANSFORM
 from a2d_core.transform.swa_attention import has_sliding_window_seam
 
@@ -29,7 +29,7 @@ def load_model(model_dir: str | Path, dtype: str = "float32") -> tuple[Any, Any]
     from a2d_core.device import select_dtype
 
     model = AutoModelForCausalLM.from_pretrained(
-        str(model_dir), attn_implementation="eager", torch_dtype=select_dtype(dtype)
+        str(model_dir), attn_implementation="eager", dtype=select_dtype(dtype)
     ).eval()
     tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
     return model, tokenizer
@@ -87,13 +87,13 @@ def resolve_capabilities(model: Any) -> list[str]:
     """
     if has_sliding_window_seam(model):
         return ["attn.swa"]
-    if _find_causal_mask_owner(model) is not None:
+    if is_rope_family(model):
         return ["attn.gqa"]
     if any(type(m).__name__ == "GPT2Attention" for m in model.modules()):
         return ["attn.full"]
     raise ValueError(
         "no supported attention seam found on model "
-        f"{type(model).__name__!r} (neither _update_causal_mask nor GPT2Attention)"
+        f"{type(model).__name__!r} (neither RoPE-family nor GPT2Attention)"
     )
 
 

packages/a2d-worker-hf/src/a2d_core/transform/attention.py

@@ -31,11 +31,16 @@
 from __future__ import annotations
 
 import math
+import sys
 from dataclasses import dataclass
 from typing import Any
 
 import torch
 
+# Prefix of the per-``AnnealState`` implementation key registered in HF's attention and
+# mask registries; also how an already-installed model is recognised on re-install.
+_KEY_PREFIX = "a2d_annealed_eager_"
+
 
 @dataclass
 class AnnealState:
@@ -59,97 +64,86 @@ def schedule(step: int, anneal_steps: int, kind: str = "linear") -> float:
     return max(0.0, min(step / anneal_steps, 1.0))
 
 
-def annealed_additive_mask(
-    q_len: int, k_len: int, alpha: float, dtype: torch.dtype, device: torch.device
-) -> torch.Tensor:
-    """Build the ``[q_len, k_len]`` additive pre-softmax mask for one attention call.
+def _reveal_penalty(alpha: float, dtype: torch.dtype) -> float:
+    """Additive pre-softmax penalty for a revealed masked cell at this ``alpha``.
 
-    ``0`` on/below the causal diagonal, ``clamp(log(alpha), finfo.min)`` strictly
-    above it. At ``alpha=0`` the future penalty is exactly ``finfo(dtype).min`` so
-    ``score + penalty`` reproduces base's ``torch.where``-to-``finfo.min`` to the bit.
+    ``finfo(dtype).min`` at ``alpha<=0`` (matches the base mask exactly, so the identity
+    gate is bit-identical); ``log(alpha)`` clamped to ``finfo.min`` otherwise; ``0`` at
+    ``alpha=1`` (fully bidirectional).
     """
     finfo_min = torch.finfo(dtype).min
-    # Match base's slice (module.bias[:, :, k-q:k, :k]) so the offset/cached q<k case
-    # stays aligned with GPT-2 causality.
-    causal = torch.tril(torch.ones(k_len, k_len, dtype=torch.bool, device=device))[
-        k_len - q_len : k_len, :k_len
-    ]
-    penalty = finfo_min if alpha <= 0.0 else max(math.log(alpha), finfo_min)
-    zero = torch.zeros((), dtype=dtype, device=device)
-    pen = torch.full((), penalty, dtype=dtype, device=device)
-    # Rebuilt per call, mirroring base's per-call causal slice.
-    return torch.where(causal, zero, pen)
-
-
-# The true library ``eager_attention_forward``, captured once when the global patch
-# is installed so the replacement can delegate to it (reusing base's exact scaling,
-# softmax and value matmul keeps the alpha=0 path bit-identical by construction).
-_original_eager: Any = None
-
-
-def _patched_eager(
-    module: Any,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: torch.Tensor | None,
-    head_mask: torch.Tensor | None = None,
-    **kwargs: Any,
-) -> Any:
-    """Drop-in for ``eager_attention_forward``: inject the annealed additive mask.
-
-    Only acts on modules tagged with ``_a2d_anneal`` (whose ``self.bias`` we have
-    neutralized to all-True); every other module delegates unchanged, so a base
-    reference model sharing this process keeps genuine causal attention.
-    """
-    state: AnnealState | None = getattr(module, "_a2d_anneal", None)
-    if state is None:
-        return _original_eager(
-            module, query, key, value, attention_mask, head_mask=head_mask, **kwargs
-        )
-    add = annealed_additive_mask(
-        query.size(-2), key.size(-2), state.alpha, query.dtype, query.device
-    )
-    if attention_mask is not None:
-        # Fold any real (padding) mask in; base adds it after its where, we add once.
-        add = add + attention_mask
-    return _original_eager(module, query, key, value, add, head_mask=head_mask, **kwargs)
+    if alpha <= 0.0:
+        return finfo_min
+    return max(math.log(alpha), finfo_min)
 
 
-def _ensure_global_patch() -> None:
-    """Install the process-global eager-attention replacement exactly once."""
-    global _original_eager
-    if _original_eager is not None:
-        return
-    import transformers.models.gpt2.modeling_gpt2 as modeling_gpt2
+def annealed_eager_mask(state: AnnealState, **kwargs: Any) -> Any:
+    """``ALL_MASK_ATTENTION_FUNCTIONS`` entry: HF's eager mask with the annealed reveal.
 
-    _original_eager = modeling_gpt2.eager_attention_forward
-    modeling_gpt2.eager_attention_forward = _patched_eager
+    Works for every mask flavour built through this interface - the causal mask, and the
+    sliding-window mask Gemma 2/3 and Mistral additionally request - because both arrive
+    here as a 4D additive mask whose masked cells are exactly ``finfo(dtype).min``.
+    """
+    from transformers.masking_utils import eager_mask
+
+    base = eager_mask(**kwargs)
+    if base is None:
+        return None
+    dtype = base.dtype
+    reveal = base == torch.finfo(dtype).min
+    # Preserve genuine padding: only reveal a masked cell whose key is a real token, so
+    # a padded key stays finfo.min at every alpha.
+    padding = kwargs.get("attention_mask")
+    if isinstance(padding, torch.Tensor) and padding.dim() == 2:
+        reveal = reveal & (padding != 0)[:, None, None, : base.shape[-1]]
+    penalty = torch.full((), _reveal_penalty(state.alpha, dtype), dtype=dtype, device=base.device)
+    return torch.where(reveal, penalty, base)
+
+
+def _a2d_eager_attention(module: Any, *args: Any, **kwargs: Any) -> Any:
+    """``ALL_ATTENTION_FUNCTIONS`` entry: the attention module's own eager forward.
+
+    Both registries are keyed by ``config._attn_implementation``, so pointing a model at
+    a custom key also moves it off HF's built-in ``"eager"`` attention path. Resolving
+    ``eager_attention_forward`` from the module that defines the attention class keeps
+    each family's own eager maths (Gemma 2 logit softcapping, GPT-2's attention scaling)
+    exactly as base computes it - only the additive mask is ours.
+    """
+    return sys.modules[type(module).__module__].eager_attention_forward(module, *args, **kwargs)
 
 
-def install_anneal_patch(model: Any, state: AnnealState) -> None:
-    """Route ``model``'s eager attention through the annealed additive mask.
+def install_mask_anneal(model: Any, state: AnnealState) -> None:
+    """Route ``model``'s mask construction through the annealed reveal.
 
-    Requires ``attn_implementation="eager"`` (Decision 2); forces ``use_cache=False``,
-    neutralizes each ``GPT2Attention``'s causal ``self.bias`` (all-True) so only the
-    additive mask governs causality, and tags each attention module with the shared
-    ``state``.
+    Requires ``attn_implementation="eager"`` (Decision 2); forces ``use_cache=False``
+    (diffusion decodes the full canvas, ARCHITECTURE.md §7). The implementation key is
+    derived from ``state``, which the registry entry keeps alive, so the key is stable
+    and unique per anneal state: re-installing the same state is idempotent, while
+    installing a fresh state swaps in a new key rather than being silently ignored.
     """
-    from transformers.models.gpt2.modeling_gpt2 import GPT2Attention
+    from transformers import AttentionInterface
+    from transformers.masking_utils import AttentionMaskInterface
 
-    if model.config._attn_implementation != "eager":
+    implementation = model.config._attn_implementation
+    if implementation != "eager" and not implementation.startswith(_KEY_PREFIX):
         raise ValueError(
             "anneal patch requires attn_implementation='eager', got "
-            f"{model.config._attn_implementation!r} (Decision 2)"
+            f"{implementation!r} (Decision 2)"
         )
+    key = f"{_KEY_PREFIX}{id(state)}"
+    AttentionMaskInterface.register(key, lambda **kwargs: annealed_eager_mask(state, **kwargs))
+    AttentionInterface.register(key, _a2d_eager_attention)
     model.config.use_cache = False
-    patched_any = False
-    for module in model.modules():
-        if isinstance(module, GPT2Attention):
-            n = module.bias.shape[-1]
-            module.bias = torch.ones((1, 1, n, n), dtype=torch.bool, device=module.bias.device)
-            module._a2d_anneal = state  # dynamic tag read back in _patched_eager
-            patched_any = True
-    if not patched_any:
+    model.set_attn_implementation(key)
+    if model.config._attn_implementation != key:
+        raise ValueError(f"transformers refused attn_implementation={key!r}: alpha stays causal")
+
+
+def install_anneal_patch(model: Any, state: AnnealState) -> None:
+    """``attn.full`` install: the shared mask anneal, gated on GPT-2's dense attention."""
+    if not any(type(m).__name__ == "GPT2Attention" for m in model.modules()):
         raise ValueError("anneal patch found no GPT2Attention modules (not a GPT-2 eager model?)")
-    _ensure_global_patch()
+    if getattr(model.config, "reorder_and_upcast_attn", False):
+        # That path survives only a literal "eager" key: base keeps it, patched would not.
+        raise ValueError("anneal patch does not support reorder_and_upcast_attn=True")
+    install_mask_anneal(model, state)

packages/a2d-worker-hf/src/a2d_core/transform/gqa_attention.py

@@ -39,104 +39,27 @@
 
 from __future__ import annotations
 
-import math
-import types
 from typing import Any
 
-import torch
+from a2d_core.transform.attention import AnnealState, install_mask_anneal
 
-from a2d_core.transform.attention import AnnealState
 
-# Sentinel stashing the original bound ``_update_causal_mask`` on the decoder module
-# once patched; also guards against double-wrapping (which would capture our own
-# wrapper as the "original" and anneal twice).
-_ORIG_ATTR = "_a2d_gqa_orig_update_causal_mask"
+def is_rope_family(model: Any) -> bool:
+    """True iff ``model`` is a RoPE-family decoder (the ``attn.gqa`` structural signal).
 
-# The live ``AnnealState`` on the decoder module, read by the wrapper on every call
-# and re-assigned by every install (parity with the GPT-2 seam's ``_a2d_anneal`` tag),
-# so re-installing with a fresh state swaps it in instead of being silently ignored.
-_STATE_ATTR = "_a2d_anneal"
-
-
-def _reveal_penalty(alpha: float, dtype: torch.dtype) -> float:
-    """Additive pre-softmax penalty for a revealed masked cell at this ``alpha``.
-
-    ``finfo(dtype).min`` at ``alpha<=0`` (matches the base causal mask exactly, so the
-    identity gate is bit-identical); ``log(alpha)`` clamped to ``finfo.min`` otherwise;
-    ``0`` at ``alpha=1`` (fully bidirectional).
-    """
-    finfo_min = torch.finfo(dtype).min
-    if alpha <= 0.0:
-        return finfo_min
-    return max(math.log(alpha), finfo_min)
-
-
-def _find_causal_mask_owner(model: Any) -> Any | None:
-    """The decoder submodule (the ``*Model``) that defines ``_update_causal_mask``.
-
-    For ``*ForCausalLM`` this is ``model.model``; found structurally (by the class
-    that actually defines the method) so no family module names are hardcoded.
+    Keyed off ``num_key_value_heads``, which every model in this family declares and
+    GPT-2 does not - the same attribute detect reads to emit the capability - so no
+    family module names are hardcoded. MQA, GQA and full-attention variants all declare
+    it, which is exactly the set this seam covers.
     """
-    for module in model.modules():
-        if type(module).__dict__.get("_update_causal_mask") is not None:
-            return module
-    return None
+    return getattr(model.config, "num_key_value_heads", None) is not None
 
 
 def install_gqa_anneal_patch(model: Any, state: AnnealState) -> None:
-    """Route ``model``'s causal-mask construction through the annealed reveal.
-
-    Requires ``attn_implementation="eager"`` (Decision 2, parity with the GPT-2 seam);
-    forces ``use_cache=False`` (diffusion decodes the full canvas, ARCHITECTURE.md §7);
-    wraps the decoder's ``_update_causal_mask`` and tags the owner with ``state``,
-    re-assigned on every install so a later install swaps in its fresh state.
-    """
-    if model.config._attn_implementation != "eager":
-        raise ValueError(
-            "anneal patch requires attn_implementation='eager', got "
-            f"{model.config._attn_implementation!r} (Decision 2)"
-        )
-    owner = _find_causal_mask_owner(model)
-    if owner is None:
+    """``attn.gqa`` install: the shared mask anneal, gated on the RoPE-family signal."""
+    if not is_rope_family(model):
         raise ValueError(
-            "GQA anneal patch found no _update_causal_mask seam (not a RoPE-family causal model?)"
+            "GQA anneal patch found no RoPE-family causal mask seam on "
+            f"{type(model).__name__!r} (no num_key_value_heads)"
         )
-    model.config.use_cache = False
-    setattr(owner, _STATE_ATTR, state)
-
-    if hasattr(owner, _ORIG_ATTR):
-        return  # already wrapped; wrapping again would double-anneal
-    orig = owner._update_causal_mask  # bound original, captured before we shadow it
-    setattr(owner, _ORIG_ATTR, orig)
-
-    def _annealed_update_causal_mask(
-        self: Any,
-        attention_mask: Any,
-        input_tensor: Any,
-        cache_position: Any,
-        *args: Any,
-        **kwargs: Any,
-    ) -> Any:
-        base_mask = orig(attention_mask, input_tensor, cache_position, *args, **kwargs)
-        # Eager always returns a 4D additive mask; None only appears on sdpa/flash
-        # fast paths we never take. Nothing to anneal if it is absent.
-        if base_mask is None:
-            return None
-        dtype = base_mask.dtype
-        kv_len = base_mask.shape[-1]
-        # Every cell the base mask masked (exactly finfo.min): the strictly-future
-        # causal cells, plus - for single-mask windowed families like Mistral v0.1 or
-        # Qwen2 with an active sliding window - the far-past out-of-window cells.
-        reveal = base_mask == torch.finfo(dtype).min
-        # Preserve genuine padding: only reveal a masked cell whose key is real, so a
-        # padded key stays finfo.min at every alpha (parity with the GPT-2 seam).
-        if isinstance(attention_mask, torch.Tensor) and attention_mask.dim() == 2:
-            real_key = (attention_mask != 0)[:, None, None, :kv_len]
-            reveal = reveal & real_key
-        live_state: AnnealState = getattr(self, _STATE_ATTR)
-        penalty = torch.full(
-            (), _reveal_penalty(live_state.alpha, dtype), dtype=dtype, device=base_mask.device
-        )
-        return torch.where(reveal, penalty, base_mask)
-
-    owner._update_causal_mask = types.MethodType(_annealed_update_causal_mask, owner)
+    install_mask_anneal(model, state)