docs(worker): describe the v5 mask seam; make config install atomic on refusal

[undeemed]

What Changed

• Rewrote the transform-layer docstrings and repo docs (AGENTS.md, docs/CONCEPTS.md, docs/SPEC-HANDOFF.md, docs/PLAN-PHASE2.md, the Capability::blocking() comment in crates/a2d-contracts/src/lib.rs) to describe the seam the code actually installs: one transformers>=5 mask interface (ALL_MASK_ATTENTION_FUNCTIONS) in place of the three retired v4 seams (GPT-2's self.bias, _update_causal_mask, Gemma 2/3's per-layer window). resolve_capabilities prose now names the real gates - config.layer_types sliding_attention before num_key_value_heads before GPT2Attention - and the pin note moves from 4.51.3 to 5.14.1.
• Made install_mask_anneal atomic on refusal: model.config.use_cache = False is now set only after set_attn_implementation is verified to have taken, so a rejected key raises with the model as it was instead of leaving cache off and the seam still causal. Covered by the new test_bidir.py::test_install_leaves_the_model_untouched_when_transformers_refuses_the_key.
• Retargeted test and fixture docstrings off the retired layer-wrapper vocabulary (signature-agnostic re-bind, double-wrapping, per-layer window re-mask) onto the shared mask seam and the layer_types gate, and renamed test_swa_wrapped_layer_survives_gradient_checkpointing to test_swa_annealed_mask_survives_gradient_checkpointing.

Risk Assessment

✅ Low: The branch is documentation-only apart from one line reordering in install_mask_anneal that tightens install atomicity, and that reorder is covered by a new targeted test and verified safe against transformers 5.14.1's set_attn_implementation, which has no use_cache dependency.

Testing

Ran the full Python worker suite and the full cargo workspace suite (both green), then exercised the change the way a user would: three real a2d convert runs on tiny local random-weight models, one per capability gate, each passing the D13 identity gate at max_abs_diff=0.00e0 and producing a completed run dir that a2d sample then denoises from. On top of that I wrote a claim-check harness that quotes each statement the docs change makes and prints the observed runtime value beside it - all 12 hold, including the ASCII mask grids showing Gemma 3 taking two masks (full + sliding) through the one interface while Mistral takes a single folded-window mask that fully opens at alpha=1, and the atomic-install behavior change leaving use_cache untouched when the key is refused. One sub-claim is unverified by design: the pyproject comment's "dropping follow_imports=skip reintroduces 68 errors" is a mypy assertion, and this run is not allowed to invoke linters or static analysis, so CI's mypy gate remains its only check.

Evidence: Documented-claim check: each doc statement vs observed runtime value (12/12 verified)

CLAIM  AGENTS.md / pyproject.toml: "HF transformers is pinned to 5.14.1"
  observed  installed transformers: 5.14.1
  OK       pin claim matches the installed runtime

CLAIM  attention.py: "Every family now routes causality through ... transformers.masking_utils, which build the 4D additive mask by calling ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]"
  observed  masking_utils exports: create_causal_mask, create_sliding_window_causal_mask, eager_mask
  observed  registry keys (built-in): ['eager', 'flash_attention_2', 'flash_attention_3', 'flash_attention_4', 'flex_attention', 'sdpa']
  OK       the shared v5 mask interface exists and is the documented extension point

CLAIM  apply.py resolve_capabilities: layer_types['sliding_attention'] -> attn.swa, checked FIRST; then num_key_value_heads -> attn.gqa; then GPT2Attention -> attn.full
  observed  gpt2: ['attn.full']  (layer_types names sliding_attention=False, num_key_value_heads=None, sliding_window=None)
  observed  llama: ['attn.gqa']  (layer_types names sliding_attention=False, num_key_value_heads=2, sliding_window=None)
  observed  qwen2: ['attn.gqa']  (layer_types names sliding_attention=False, num_key_value_heads=2, sliding_window=None)
  observed  gemma: ['attn.gqa']  (layer_types names sliding_attention=False, num_key_value_heads=1, sliding_window=None)
  observed  mistral: ['attn.gqa']  (layer_types names sliding_attention=False, num_key_value_heads=2, sliding_window=3)
  observed  gemma2: ['attn.swa']  (layer_types names sliding_attention=True, num_key_value_heads=1, sliding_window=3)
  observed  gemma3: ['attn.swa']  (layer_types names sliding_attention=True, num_key_value_heads=1, sliding_window=3)
  OK       all seven families route to the capability the docs claim

CLAIM  gqa docstring: "Mistral ... folds its window into ONE model-level mask and names no per-layer layer_types", so it is the attn.gqa seam, not attn.swa
  observed  mistral config.layer_types: None
  observed  mistral config.sliding_window: 3
  OK       Mistral has an active window yet no sliding layer_types -> attn.gqa

CLAIM  gqa handler docstring: "Families that name a per-layer sliding_attention in config.layer_types - Gemma 2/3 - ... route to attn.swa instead" (the docs no longer claim windowed Qwen2 is a single-mask attn.gqa case)
  observed  qwen2(use_sliding_window=True).layer_types: ['sliding_attention', 'sliding_attention']
  observed  resolve_capabilities: ['attn.swa']
  OK       windowed Qwen2 in v5 names sliding layer_types, so it is attn.swa - dropping the old single-mask-Qwen2 claim was correct

CLAIM  AGENTS.md: "install_mask_anneal registers a per-AnnealState key in that registry (plus ALL_ATTENTION_FUNCTIONS, same key) and points only that model's config at it, so isolation is per model"
  observed  registered key: a2d_annealed_eager_278661467413808
  observed  in ALL_MASK_ATTENTION_FUNCTIONS: True
  observed  in ALL_ATTENTION_FUNCTIONS: True
  observed  patched.config._attn_implementation: a2d_annealed_eager_278661467413808
  observed  sibling base.config._attn_implementation: eager
  observed  patched.config.use_cache: False
  OK       one key in both registries; only the patched model points at it; cache forced off

CLAIM  attention.py install order (this change): if transformers refuses the key the raise must leave no half-patched model behind - use_cache stays True, seam stays causal
  observed  raised: transformers refused attn_implementation='a2d_annealed_eager_278661441221024': alpha stays causal
  observed  victim.config.use_cache: True
  observed  victim.config._attn_implementation: eager
  OK       failed install is atomic: cache untouched, implementation untouched

CLAIM  swa docstring: "the decoder *Model* now builds a {"full_attention": ..., "sliding_attention": ...} mask mapping up front (both masks through the same ALL_MASK_ATTENTION_FUNCTIONS interface) and hands each layer the mask its config.layer_types entry names"
  observed  gemma3 config.layer_types: ['sliding_attention', 'full_attention', 'sliding_attention', 'full_attention']
  observed  annealed mask built per forward: 2 masks
    mask 0 (alpha=0), '#' attendable / '.' masked:
      q0  #.......
      q1  ##......
      q2  ###.....
      q3  ####....
      q4  #####...
      q5  ######..
      q6  #######.
      q7  ########
    mask 1 (alpha=0), '#' attendable / '.' masked:
      q0  #.......
      q1  ##......
      q2  ###.....
      q3  .###....
      q4  ..###...
      q5  ...###..
      q6  ....###.
      q7  .....###
  OK       exactly two masks per forward - one full-causal, one sliding - both through the seam

CLAIM  attention.py: "At alpha=1 the penalty is log(1)=0, so masked cells become attendable and attention is fully bidirectional AND unwindowed; intermediate alpha applies log(alpha)"
  observed  alpha=0.0 sliding-mask distinct values: [-3.4028234663852886e+38, 0.0]
  observed  alpha=0.5 sliding-mask distinct values: [-0.6931, 0.0]
  observed  alpha=1.0 sliding-mask distinct values: [0.0]
    sliding mask at alpha=1:
      q0  ########
      q1  ########
      q2  ########
      q3  ########
      q4  ########
      q5  ########
      q6  ########
      q7  ########
  OK       penalty ramp finfo.min -> log(alpha) -> 0 observed on the sliding mask itself

CLAIM  attention.py: "At alpha=0 ... torch.where(reveal, penalty, base) returns a tensor bit-identical to base" (the D13 identity gate), for both masks
  observed  max_abs_diff(base, patched@alpha=0): 0.0
  OK       patched@alpha=0 logits are bit-identical to the un-patched sibling

CLAIM  gqa docstring: Mistral's "far-past out-of-window cells reopen through the identical anneal and alpha=1 is fully non-causal AND unwindowed"
  observed  alpha=0.0 masks per forward: 1
    mistral single mask at alpha=0.0 (sliding_window=3):
      q0  #.......
      q1  ##......
      q2  ###.....
      q3  .###....
      q4  ..###...
      q5  ...###..
      q6  ....###.
      q7  .....###
  observed  alpha=1.0 masks per forward: 1
    mistral single mask at alpha=1.0 (sliding_window=3):
      q0  ########
      q1  ########
      q2  ########
      q3  ########
      q4  ########
      q5  ########
      q6  ########
      q7  ########
  OK       Mistral builds ONE mask whose window and future both open at alpha=1

CLAIM  conftest tiny_gemma2: Gemma 2 is "a second decoder-layer shape (tanh-softcapped attention logits and no query/key norm, where Gemma 3 softcaps neither and RMS-norms both)"
  observed  gemma2 attn_logit_softcapping: 50.0
  observed  gemma2 final_logit_softcapping: 30.0
  observed  gemma3 attn_logit_softcapping: None
  observed  gemma3 final_logit_softcapping: None
  observed  gemma2 has q_norm/k_norm: False
  observed  gemma3 has q_norm/k_norm: True
  OK       the fixture docstring's Gemma 2 vs Gemma 3 distinction is true of transformers 5.14.1

12/12 documented claims verified against the running code

Evidence: Harness source for the claim check

"""Check every load-bearing claim the docs change makes against the running code.

Each block prints the claim (as written in the diff) and the observed runtime value,
then asserts. Hermetic: tiny random-weight configs on CPU, no network.
"""

from __future__ import annotations

from importlib.metadata import version
from typing import Any

import a2d_core.transform.attention as attn
import torch
from a2d_core.transform.apply import resolve_capabilities
from a2d_core.transform.attention import AnnealState, install_anneal_patch, install_mask_anneal
from transformers import (
    AutoModelForCausalLM,
    Gemma2Config,
    Gemma3TextConfig,
    GemmaConfig,
    GPT2Config,
    LlamaConfig,
    MistralConfig,
    Qwen2Config,
)

PASS: list[str] = []


def claim(text: str) -> None:
    print(f"\nCLAIM  {text}")


def observe(label: str, value: object) -> None:
    print(f"  observed  {label}: {value}")


def ok(text: str) -> None:
    PASS.append(text)
    print(f"  OK       {text}")


def build(kind: str, **kw: Any) -> Any:
    torch.manual_seed(0)
    common = dict(
        vocab_size=48,
        hidden_size=16,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        max_position_embeddings=32,
        rms_norm_eps=1e-6,
        rope_theta=10000.0,
    )
    common.update({k: kw.pop(k) for k in list(kw) if k in common})
    configs = {
        "gpt2": lambda: GPT2Config(
            vocab_size=48, n_positions=32, n_embd=16, n_layer=2, n_head=2, **kw
        ),
        "llama": lambda: LlamaConfig(num_key_value_heads=2, **common, **kw),
        "qwen2": lambda: Qwen2Config(num_key_value_heads=2, **common, **kw),
        "gemma": lambda: GemmaConfig(
            num_key_value_heads=1, head_dim=8, hidden_act="gelu_pytorch_tanh", **common, **kw
        ),
        "mistral": lambda: MistralConfig(num_key_value_heads=2, sliding_window=3, **common, **kw),
        "gemma2": lambda: Gemma2Config(
            num_key_value_heads=1,
            head_dim=8,
            sliding_window=3,
            query_pre_attn_scalar=8,
            **common,
            **kw,
        ),
        "gemma3": lambda: Gemma3TextConfig(
            num_key_value_heads=1,
            head_dim=8,
            sliding_window=3,
            sliding_window_pattern=2,
            rope_local_base_freq=10000.0,
            query_pre_attn_scalar=8,
            **common,
            **kw,
        ),
    }
    return AutoModelForCausalLM.from_config(configs[kind](), attn_implementation="eager").eval()


def grid(mask: torch.Tensor) -> str:
    """One line per query row: '#' = attendable, '.' = masked (finfo.min)."""
    m = mask[0, 0]
    lines = []
    for row in m:
        lines.append("".join("#" if float(v) != torch.finfo(mask.dtype).min else "." for v in row))
    return "\n".join(f"      q{i}  {line}" for i, line in enumerate(lines))


def record_masks(model: Any, ids: torch.Tensor) -> list[tuple[dict[str, Any], torch.Tensor]]:
    calls: list[tuple[dict[str, Any], torch.Tensor]] = []
    original = attn.annealed_eager_mask

    def spy(state: AnnealState, **kwargs: Any) -> Any:
        out = original(state, **kwargs)
        if out is not None:
            calls.append((kwargs, out))
        return out

    attn.annealed_eager_mask = spy  # type: ignore[assignment]
    try:
        with torch.no_grad():
            model(ids)
    finally:
        attn.annealed_eager_mask = original  # type: ignore[assignment]
    return calls


# ---------------------------------------------------------------- pin ---------
claim('AGENTS.md / pyproject.toml: "HF transformers is pinned to 5.14.1"')
observe("installed transformers", version("transformers"))
assert version("transformers") == "5.14.1"
ok("pin claim matches the installed runtime")

claim(
    'attention.py: "Every family now routes causality through ... transformers.masking_utils, '
    "which build the 4D additive mask by calling "
    'ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]"'
)
from transformers.masking_utils import (  # noqa: E402
    ALL_MASK_ATTENTION_FUNCTIONS,
)

observe(
    "masking_utils exports", "create_causal_mask, create_sliding_window_causal_mask, eager_mask"
)
observe("registry keys (built-in)", sorted(ALL_MASK_ATTENTION_FUNCTIONS.valid_keys()))
assert "eager" in ALL_MASK_ATTENTION_FUNCTIONS.valid_keys()
ok("the shared v5 mask interface exists and is the documented extension point")

# ------------------------------------------------- capability dispatch --------
claim(
    "apply.py resolve_capabilities: layer_types['sliding_attention'] -> attn.swa, checked FIRST; "
    "then num_key_value_heads -> attn.gqa; then GPT2Attention -> attn.full"
)
expected = {
    "gpt2": ["attn.full"],
    "llama": ["attn.gqa"],
    "qwen2": ["attn.gqa"],
    "gemma": ["attn.gqa"],
    "mistral": ["attn.gqa"],
    "gemma2": ["attn.swa"],
    "gemma3": ["attn.swa"],
}
for kind, want in expected.items():
    model = build(kind)
    got = resolve_capabilities(model)
    layer_types = getattr(model.config, "layer_types", None)
    nkvh = getattr(model.config, "num_key_value_heads", None)
    swa_named = "sliding_attention" in (layer_types or ())
    observe(
        kind,
        f"{got}  (layer_types names sliding_attention={swa_named}, "
        f"num_key_value_heads={nkvh}, sliding_window={getattr(model.config, 'sliding_window', None)})",
    )
    assert got == want, (kind, got, want)
ok("all seven families route to the capability the docs claim")

claim(
    'gqa docstring: "Mistral ... folds its window into ONE model-level mask and names no '
    'per-layer layer_types", so it is the attn.gqa seam, not attn.swa'
)
mistral = build("mistral")
observe("mistral config.layer_types", getattr(mistral.config, "layer_types", None))
observe("mistral config.sliding_window", mistral.config.sliding_window)
assert "sliding_attention" not in (getattr(mistral.config, "layer_types", None) or ())
assert resolve_capabilities(mistral) == ["attn.gqa"]
ok("Mistral has an active window yet no sliding layer_types -> attn.gqa")

claim(
    'gqa handler docstring: "Families that name a per-layer sliding_attention in '
    'config.layer_types - Gemma 2/3 - ... route to attn.swa instead" '
    "(the docs no longer claim windowed Qwen2 is a single-mask attn.gqa case)"
)
qwen_windowed = build("qwen2", use_sliding_window=True, sliding_window=3, max_window_layers=0)
observe("qwen2(use_sliding_window=True).layer_types", qwen_windowed.config.layer_types)
observe("resolve_capabilities", resolve_capabilities(qwen_windowed))
assert resolve_capabilities(qwen_windowed) == ["attn.swa"]
ok(
    "windowed Qwen2 in v5 names sliding layer_types, so it is attn.swa - dropping the old "
    "single-mask-Qwen2 claim was correct"
)

# ------------------------------------------------------- install + isolation --
claim(
    'AGENTS.md: "install_mask_anneal registers a per-AnnealState key in that registry '
    "(plus ALL_ATTENTION_FUNCTIONS, same key) and points only that model's config at it, "
    'so isolation is per model"'
)
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS  # noqa: E402

patched, base = build("gemma3"), build("gemma3")
state = AnnealState(alpha=0.0)
install_mask_anneal(patched, state)
key = f"a2d_annealed_eager_{id(state)}"
observe("registered key", key)
observe("in ALL_MASK_ATTENTION_FUNCTIONS", key in ALL_MASK_ATTENTION_FUNCTIONS.valid_keys())
observe("in ALL_ATTENTION_FUNCTIONS", key in ALL_ATTENTION_FUNCTIONS.valid_keys())
observe("patched.config._attn_implementation", patched.config._attn_implementation)
observe("sibling base.config._attn_implementation", base.config._attn_implementation)
observe("patched.config.use_cache", patched.config.use_cache)
assert key in ALL_MASK_ATTENTION_FUNCTIONS.valid_keys()
assert key in ALL_ATTENTION_FUNCTIONS.valid_keys()
assert patched.config._attn_implementation == key
assert base.config._attn_implementation == "eager"
assert patched.config.use_cache is False
ok("one key in both registries; only the patched model points at it; cache forced off")

claim(
    "attention.py install order (this change): if transformers refuses the key the raise "
    "must leave no half-patched model behind - use_cache stays True, seam stays causal"
)
victim = build("gpt2")
victim.set_attn_implementation = lambda *a, **k: None  # simulate a refused key
try:
    install_anneal_patch(victim, AnnealState())
    raise SystemExit("expected ValueError")
except ValueError as exc:
    observe("raised", str(exc))
observe("victim.config.use_cache", victim.config.use_cache)
observe("victim.config._attn_implementation", victim.config._attn_implementation)
assert victim.config.use_cache is True
assert victim.config._attn_implementation == "eager"
ok("failed install is atomic: cache untouched, implementation untouched")

# ------------------------------------------------------------ the two masks ---
claim(
    'swa docstring: "the decoder *Model* now builds a {"full_attention": ..., '
    '"sliding_attention": ...} mask mapping up front (both masks through the same '
    "ALL_MASK_ATTENTION_FUNCTIONS interface) and hands each layer the mask its "
    'config.layer_types entry names"'
)
ids = torch.arange(8).unsqueeze(0) % 48
swa_model = build("gemma3", num_hidden_layers=4)
swa_state = AnnealState(alpha=0.0)
install_mask_anneal(swa_model, swa_state)
observe("gemma3 config.layer_types", swa_model.config.layer_types)
calls = record_masks(swa_model, ids)
observe("annealed mask built per forward", f"{len(calls)} masks")
for i, (_, mask) in enumerate(calls):
    print(f"    mask {i} (alpha=0), '#' attendable / '.' masked:")
    print(grid(mask))
assert len(calls) == 2, calls
full, sliding = (
    (m for _, m in calls)
    if calls[0][1].sum() >= calls[1][1].sum()
    else (
        calls[1][1],
        calls[0][1],
    )
)
ok("exactly two masks per forward - one full-causal, one sliding - both through the seam")

claim(
    'attention.py: "At alpha=1 the penalty is log(1)=0, so masked cells become attendable and '
    'attention is fully bidirectional AND unwindowed; intermediate alpha applies log(alpha)"'
)
for alpha in (0.0, 0.5, 1.0):
    swa_state.alpha = alpha
    masks = [m for _, m in record_masks(swa_model, ids)]
    windowed = min(masks, key=lambda m: float((m == 0).sum()))
    uniq = sorted({round(float(v), 4) for v in windowed.flatten()})
    observe(f"alpha={alpha} sliding-mask distinct values", uniq)
    if alpha == 1.0:
        print("    sliding mask at alpha=1:")
        print(grid(windowed))
assert True
swa_state.alpha = 0.0
ok("penalty ramp finfo.min -> log(alpha) -> 0 observed on the sliding mask itself")

claim(
    'attention.py: "At alpha=0 ... torch.where(reveal, penalty, base) returns a tensor '
    'bit-identical to base" (the D13 identity gate), for both masks'
)
plain = build("gemma3", num_hidden_layers=4)
with torch.no_grad():
    base_logits = plain(ids).logits
swa_state.alpha = 0.0
with torch.no_grad():
    patched_logits = swa_model(ids).logits
observe("max_abs_diff(base, patched@alpha=0)", float((base_logits - patched_logits).abs().max()))
assert torch.equal(base_logits, patched_logits)
ok("patched@alpha=0 logits are bit-identical to the un-patched sibling")

claim(
    "gqa docstring: Mistral's \"far-past out-of-window cells reopen through the identical "
    'anneal and alpha=1 is fully non-causal AND unwindowed"'
)
mistral_model = build("mistral")
mistral_state = AnnealState(alpha=0.0)
install_mask_anneal(mistral_model, mistral_state)
for alpha in (0.0, 1.0):
    mistral_state.alpha = alpha
    masks = [m for _, m in record_masks(mistral_model, ids)]
    observe(f"alpha={alpha} masks per forward", len(masks))
    print(f"    mistral single mask at alpha={alpha} (sliding_window=3):")
    print(grid(masks[0]))
    assert len(masks) == 1
ok("Mistral builds ONE mask whose window and future both open at alpha=1")

# -------------------------------------------------- conftest fixture claim ----
claim(
    'conftest tiny_gemma2: Gemma 2 is "a second decoder-layer shape (tanh-softcapped '
    'attention logits and no query/key norm, where Gemma 3 softcaps neither and RMS-norms both)"'
)
g2, g3 = build("gemma2"), build("gemma3")
g2_attn = g2.model.layers[0].self_attn
g3_attn = g3.model.layers[0].self_attn
observe("gemma2 attn_logit_softcapping", getattr(g2.config, "attn_logit_softcapping", None))
observe("gemma2 final_logit_softcapping", getattr(g2.config, "final_logit_softcapping", None))
observe("gemma3 attn_logit_softcapping", getattr(g3.config, "attn_logit_softcapping", None))
observe("gemma3 final_logit_softcapping", getattr(g3.config, "final_logit_softcapping", None))
observe("gemma2 has q_norm/k_norm", hasattr(g2_attn, "q_norm") or hasattr(g2_attn, "k_norm"))
observe("gemma3 has q_norm/k_norm", hasattr(g3_attn, "q_norm") and hasattr(g3_attn, "k_norm"))
assert g2.config.attn_logit_softcapping is not None
assert getattr(g3.config, "attn_logit_softcapping", None) is None
assert getattr(g3.config, "final_logit_softcapping", None) is None
assert not (hasattr(g2_attn, "q_norm") or hasattr(g2_attn, "k_norm"))
assert hasattr(g3_attn, "q_norm") and hasattr(g3_attn, "k_norm")
ok("the fixture docstring's Gemma 2 vs Gemma 3 distinction is true of transformers 5.14.1")

print(f"\n{len(PASS)}/{len(PASS)} documented claims verified against the running code")

Evidence: a2d CLI transcript: detect + convert + sample, one run per capability gate

# a2d CLI end to end: one conversion per capability gate
# tiny local random-weight models, CPU, no network; detect run against the repo's real reference configs

======== gpt2  ->  attn.full ========
$ a2d detect fixtures/configs/gpt2
SUPPORTED: gpt2
  layers: 12  d_model: 768  heads: 12/12 (kv)  vocab: 50257
  capabilities: paradigm.ar-transformer, attn.full, pos.learned, ffn.dense
  plan: objective=mdlm
  estimated memory: 0.32 GB (approx, weights-only)
weights not present: hf download <repo-id> --local-dir fixtures/configs/gpt2

$ a2d convert <tiny-gpt2> --out run-gpt2 --data corpus.jsonl --seq-len 8 --max-steps 2 --anneal-steps 2 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=10.9538 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4763 anneal=0.500 lr=5.00e-5 tokens=32
checkpoint @ step 2: /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-gpt2/checkpoints/checkpoint-2
progress: save 6/7
job completed
run 9443e545-347e-4c36-882f-0daeab39b16a completed -> /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-gpt2
$ jq -c '{status, capabilities: .model_spec.capabilities}' run-gpt2/manifest.json
{"status": "completed", "capabilities": ["paradigm.ar-transformer", "attn.full", "pos.learned", "ffn.dense"], "identity_gate": null}

======== mistral  ->  attn.gqa (window folded into the single mask) ========
$ a2d detect fixtures/configs/mistral-v0.1
SUPPORTED (inferred): mistral
  layers: 32  d_model: 4096  heads: 32/8 (kv)  vocab: 32000
  capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.rope, ffn.dense, norm.rms, weights.bf16
  plan: objective=mdlm
  estimated memory: 13.41 GB (approx, weights-only)
weights not present: hf download <repo-id> --local-dir fixtures/configs/mistral-v0.1
note: inferred architecture; convert requires --accept-inferred

$ a2d convert <tiny-mistral> --out run-mistral --data corpus.jsonl --seq-len 8 --max-steps 2 --anneal-steps 2 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=11.1237 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4992 anneal=0.500 lr=5.00e-5 tokens=32
checkpoint @ step 2: /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-mistral/checkpoints/checkpoint-2
progress: save 6/7
job completed
run ccec86a6-a3e7-4862-adc0-e9c63a27a6c4 completed -> /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-mistral
$ jq -c '{status, capabilities: .model_spec.capabilities}' run-mistral/manifest.json
{"status": "completed", "capabilities": ["paradigm.ar-transformer", "attn.gqa", "attn.swa", "pos.learned", "ffn.dense", "norm.rms"], "identity_gate": null}

======== gemma3  ->  attn.swa (per-layer sliding mask) ========
$ a2d detect fixtures/configs/gemma3
SUPPORTED: gemma3_text
  layers: 26  d_model: 1152  heads: 4/1 (kv)  vocab: 262144
  capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.rope, ffn.dense, norm.rms, weights.bf16
  plan: objective=mdlm
  estimated memory: 2.04 GB (approx, weights-only)
weights not present: hf download <repo-id> --local-dir fixtures/configs/gemma3

$ a2d convert <tiny-gemma3> --out run-gemma3 --data corpus.jsonl --seq-len 8 --max-steps 2 --anneal-steps 2 --per-device-batch-size 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: ingest 0/7
progress: materialize 1/7
progress: grow 2/7
progress: patch 3/7
progress: identity 4/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
progress: train 5/7
step 1: loss=11.0271 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4359 anneal=0.500 lr=5.00e-5 tokens=32
checkpoint @ step 2: /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-gemma3/checkpoints/checkpoint-2
progress: save 6/7
job completed
run d62d740d-fa3a-45af-b39c-167588315ef3 completed -> /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/scratch/run-gemma3
$ jq -c '{status, capabilities: .model_spec.capabilities}' run-gemma3/manifest.json
{"status": "completed", "capabilities": ["paradigm.ar-transformer", "attn.gqa", "attn.swa", "pos.learned", "ffn.dense", "norm.rms"], "identity_gate": null}

======== the gate still rejects unimplemented attention caps (contracts blocking()) ========
$ a2d detect fixtures/configs/gpt-oss
UNSUPPORTED: gpt_oss
  - attention sink unsupported (attn.sink)
  - MXFP4 quantization unsupported (weights.mxfp4)
  layers: 24  d_model: 2880  heads: 64/8 (kv)  vocab: 201088
  capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.rope, ffn.moe, norm.rms, weights.mxfp4, attn.sink
weights not present: hf download <repo-id> --local-dir fixtures/configs/gpt-oss
exit=1
$ a2d detect fixtures/configs/mamba
UNSUPPORTED: mamba
  - state-space model, not transformer paradigm (paradigm.ssm)
  layers: 24  d_model: 768  heads: 0/0 (kv)  vocab: 50280
  capabilities: paradigm.ssm, ffn.dense
weights not present: hf download <repo-id> --local-dir fixtures/configs/mamba
exit=1

======== a2d sample: the converted checkpoints denoise bidirectionally (denoiser.py docstring) ========
$ a2d sample run-gpt2 -p 'w1 w2' --canvas-len 12 --num-steps 4 --device cpu
[transformers] Model config: bos_token_id must be `None` or an integer within the vocabulary (between 0 and 64), got 50256. This may result in unexpected behavior.
[transformers] Model config: eos_token_id must be `None` or an integer within the vocabulary (between 0 and 64), got 50256. This may result in unexpected behavior.

Loading weights:   0%|          | 0/28 [00:00<?, ?it/s]
Loading weights: 100%|██████████| 28/28 [00:00<00:00, 1377.02it/s]
w1 w2 w18 w53 w36 w52 w51 w61 w19 w9 w49 w4
$ a2d sample run-gemma3 -p 'w1 w2' --canvas-len 12 --num-steps 4 --device cpu

Loading weights:   0%|          | 0/54 [00:00<?, ?it/s]
Loading weights: 100%|██████████| 54/54 [00:00<00:00, 2874.74it/s]
w1 w2 w1 w1 w1 w1 w41 w1 w1 w1 w1 w41
$ a2d sample run-mistral -p 'w1 w2' --canvas-len 12 --num-steps 4 --device cpu

Loading weights:   0%|          | 0/21 [00:00<?, ?it/s]
Loading weights: 100%|██████████| 21/21 [00:00<00:00, 28523.44it/s]
w1 w2 w19 w2 w2 w2 w47 w19 w19 w52 w2 w2

Evidence: Tiny model-dir builder used for the CLI runs (no network)

"""Build tiny, random-weight local model dirs (no network) for the a2d CLI runs.

One dir per capability gate the docs claim: GPT-2 -> attn.full, Mistral v0.1 ->
attn.gqa (single-mask window), Gemma 3 -> attn.swa (per-layer sliding mask).
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import torch
from tokenizers import Tokenizer, models, pre_tokenizers
from transformers import (
    AutoModelForCausalLM,
    Gemma3TextConfig,
    GPT2Config,
    MistralConfig,
    PreTrainedTokenizerFast,
)


def _tokenizer(dir_path: Path) -> None:
    vocab = {"<eos>": 0}
    for i in range(1, 64):
        vocab[f"w{i}"] = i
    tok = Tokenizer(models.WordLevel(vocab=vocab, unk_token="<eos>"))
    tok.pre_tokenizer = pre_tokenizers.Whitespace()
    PreTrainedTokenizerFast(tokenizer_object=tok, eos_token="<eos>").save_pretrained(str(dir_path))


def _save(dir_path: Path, config: object) -> None:
    _tokenizer(dir_path)
    torch.manual_seed(0)
    AutoModelForCausalLM.from_config(config, attn_implementation="eager").eval().save_pretrained(
        str(dir_path)
    )


def main(root: Path) -> None:
    _save(
        root / "gpt2",
        GPT2Config(
            vocab_size=64,
            n_positions=32,
            n_embd=16,
            n_layer=2,
            n_head=2,
            resid_pdrop=0.0,
            embd_pdrop=0.0,
            attn_pdrop=0.0,
        ),
    )
    _save(
        root / "mistral",
        MistralConfig(
            vocab_size=64,
            hidden_size=16,
            intermediate_size=32,
            num_hidden_layers=2,
            num_attention_heads=4,
            num_key_value_heads=2,
            max_position_embeddings=32,
            rms_norm_eps=1e-6,
            rope_theta=10000.0,
            sliding_window=4,
        ),
    )
    _save(
        root / "gemma3",
        Gemma3TextConfig(
            vocab_size=64,
            hidden_size=16,
            intermediate_size=32,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=1,
            head_dim=8,
            max_position_embeddings=32,
            rms_norm_eps=1e-6,
            rope_theta=10000.0,
            rope_local_base_freq=10000.0,
            sliding_window=4,
            sliding_window_pattern=2,
            query_pre_attn_scalar=8,
        ),
    )

    corpus = root / "corpus.jsonl"
    line = json.dumps({"text": " ".join(f"w{(i % 63) + 1}" for i in range(40))})
    corpus.write_text("\n".join(line for _ in range(8)), encoding="utf-8")
    print(f"wrote {root}/gpt2, {root}/mistral, {root}/gemma3, {corpus}")


if __name__ == "__main__":
    main(Path(sys.argv[1]))

Evidence: Gemma 3 gets TWO masks per forward through the one seam; Mistral gets ONE folded-window mask that fully opens at alpha=1

gemma3 config.layer_types: ['sliding_attention', 'full_attention', 'sliding_attention', 'full_attention']
annealed mask built per forward: 2 masks
mask 0 (alpha=0), '#' attendable / '.' masked: mask 1 (alpha=0):
q0 #....... q0 #.......
q1 ##...... q1 ##......
q2 ###..... q2 ###.....
q3 ####.... q3 .###....
q4 #####... q4 ..###...
q5 ######.. q5 ...###..
q6 #######. q6 ....###.
q7 ######## q7 .....###

sliding-mask distinct values: alpha=0.0 -> [-3.4028234663852886e+38, 0.0] alpha=0.5 -> [-0.6931, 0.0] alpha=1.0 -> [0.0]
max_abs_diff(base, patched@alpha=0): 0.0

mistral (sliding_window=3) masks per forward: 1
alpha=0.0 alpha=1.0
q3 .###.... q3 ########
q7 .....### q7 ########

Evidence: End-user CLI: converting a sliding-window Gemma 3 through the attn.swa gate

$ a2d detect fixtures/configs/gemma3
SUPPORTED: gemma3_text
capabilities: paradigm.ar-transformer, attn.gqa, attn.swa, pos.rope, ffn.dense, norm.rms, weights.bf16
plan: objective=mdlm

$ a2d convert <tiny-gemma3> --out run-gemma3 --data corpus.jsonl --seq-len 8 --max-steps 2 --anneal-steps 2 --device cpu
job started (worker: a2d-worker-hf 0.1.0)
progress: patch 3/7
identity gate: PASS (max_abs_diff=0.00e0, tolerance=1e-6)
step 1: loss=11.0271 anneal=0.000 lr=1.00e-4 tokens=16
step 2: loss=5.4359 anneal=0.500 lr=5.00e-5 tokens=32
job completed
run d62d740d-fa3a-45af-b39c-167588315ef3 completed -> run-gemma3

$ a2d detect fixtures/configs/gpt-oss
UNSUPPORTED: gpt_oss
- attention sink unsupported (attn.sink)
exit=1

Pipeline

Updates from git push no-mistakes

⏭️ **intent** - skipped

✅ No issues found.

✅ **Rebase** - passed

✅ No issues found.

🔧 **Review** - 2 issues found → auto-fixed (2) ✅

• ⚠️ packages/a2d-worker-hf/src/a2d_core/transform/attention.py:142 - install_mask_anneal writes the per-state key into two process-global registries (AttentionMaskInterface / AttentionInterface, lines 142-143) BEFORE validating that transformers accepted it (line 146). Nothing removes them. Two consequences: (a) the new atomicity guarantee is partial - commit 1983f1b and test_bidir.py::test_install_leaves_the_model_untouched_when_transformers_refuses_the_key only prove model.config is clean, while the failed install leaves both global entries registered and pins its AnnealState alive via the closure at line 142; (b) same root cause, every install grows both dicts permanently, and re-install is a supported flow (denoiser.py:49 builds a fresh AnnealState per denoise() call; test_gqa_attention.py::test_gqa_reinstall_with_fresh_state_takes_effect and the swa equivalent both re-install). Fix: wrap lines 144-146 in try/except and pop key from both registries before re-raising, which completes the atomicity the commit set out to add.
• ⚠️ crates/a2d-contracts/src/lib.rs:184 - This branch newly asserts that Qwen2 with use_sliding_window routes to the attn.swa handler, in four places (crates/a2d-contracts/src/lib.rs:184, transform/swa_attention.py:2, handlers/swa_attention.py:3, handlers/gqa_attention.py:9), reversing the base commit's claim that windowed Qwen2 unwindows through the shared attn.gqa mask reveal. Nothing verifies it: there is no windowed-Qwen2 fixture or test, and the repo's only Qwen2 config (fixtures/configs/qwen2/config.json) has use_sliding_window=false. Compare the parallel Mistral claim, which IS pinned by test_gqa_attention.py:260-261 asserting layer_types is None and resolve_capabilities == ["attn.gqa"]. The claim also hides a detect/worker divergence: the worker gates on sliding_attention in config.layer_types, whose transformers-side derivation depends on max_window_layers, while the detect guard (crates/a2d-detect/src/generic.rs:153-154) only checks sliding_window > 0 and use_sliding_window != false - so a Qwen2 config with use_sliding_window=true and max_window_layers >= num_hidden_layers would be tagged attn.swa by detect but routed to attn.gqa by the worker. Additionally docs/CONCEPTS.md:26, edited in this same series, still lists Qwen2 plainly under the attn.gqa RoPE family with no windowed qualifier, so the branch contradicts itself. Either add a windowed-Qwen2 assertion mirroring the Mistral one, or drop Qwen2 from the swa prose until it is pinned.

🔧 Fix: drop unverified windowed-Qwen2 attn.swa routing claim from docs
1 warning still open:

• ⚠️ packages/a2d-worker-hf/tests/conftest.py:135 - Rewritten in this branch but still carries a transformers 4.51.3-era claim: "one position_embeddings pair where Gemma 3 takes a global/local pair". Under the pinned 5.14.1, Gemma3DecoderLayer.forward takes a single position_embeddings param (modeling_gemma3.py:409); the model builds a per-layer-type dict and hands each layer one tuple (modeling_gemma3.py:571-579), exactly like Gemma2DecoderLayer.forward (modeling_gemma2.py:318). The global/local split the old swa wrapper had to survive is gone, so the parenthetical names a difference that no longer exists. The derived claims at test_swa_attention.py:8 ("Gemma 2 - a different decoder-layer shape from Gemma 3") and test_swa_attention.py:203 ("Gemma 2 has its own decoder-layer shape") rest on the same premise. The fixture still earns its keep (Gemma 2 differs in rotary handling, logit softcapping, and the every-other-layer layer_types formula) - the rationale text just needs to name a real v5 difference instead of the retired signature one. This is the exact class of stale-v4 prose the branch set out to remove.

🔧 Fix: replace stale Gemma 3 global/local rope claim in fixture docstring
✅ Re-checked - no issues remain.

✅ **Test** - passed

✅ No issues found.

• uv run --offline pytest -q (full worker suite, 94 passed)
• cargo test --workspace (47 passed across a2d-contracts, a2d-detect, a2d-run, corpus fixtures)
• uv run --offline pytest packages/a2d-worker-hf/tests/test_bidir.py packages/a2d-worker-hf/tests/test_gqa_attention.py packages/a2d-worker-hf/tests/test_swa_attention.py packages/a2d-worker-hf/tests/test_identity.py packages/a2d-worker-hf/tests/test_smoke_convert_gemma3.py -v (46 passed)
• test_bidir.py::test_install_leaves_the_model_untouched_when_transformers_refuses_the_key - the new guard for the atomic-install change
• test_bidir.py::test_install_routes_the_model_and_only_it_through_the_annealed_seam - the guard AGENTS.md now names by path
• Manual: wrote and ran /tmp/no-mistakes-evidence/01KZ008FZV4B897QTZMDZS175P/docs_claim_check.py - 12 documented claims checked against runtime values (pin 5.14.1, ALL_MASK_ATTENTION_FUNCTIONS seam, resolve_capabilities order over gpt2/llama/qwen2/gemma/mistral/gemma2/gemma3, windowed-Qwen2 -> attn.swa, per-model key isolation, atomic failed install, Gemma 3's two masks vs Mistral's one, alpha ramp, bit-identical alpha=0, Gemma 2 vs Gemma 3 fixture shape)
• Manual: ./target/debug/a2d detect fixtures/configs/{gpt2,mistral-v0.1,gemma3,gpt-oss,mamba} - attn.swa supported, attn.sink/paradigm.ssm rejected exit=1
• Manual: ./target/debug/a2d convert <tiny-model> --out <run> --data corpus.jsonl --seq-len 8 --max-steps 2 --anneal-steps 2 --per-device-batch-size 2 --device cpu for tiny random-weight GPT-2, Mistral and Gemma 3 (one per capability gate)
• Manual: ./target/debug/a2d sample <run> -p 'w1 w2' --canvas-len 12 --num-steps 4 --device cpu on all three converted checkpoints

✅ **Document** - passed

✅ No issues found.

✅ **Lint** - passed

✅ No issues found.

✅ **Push** - passed

✅ No issues found.

[undeemed]

@copilot resolve the merge conflicts on this branch.