diff --git a/giant/config.py b/giant/config.py
index 93972e6..4357c70 100644
--- a/giant/config.py
+++ b/giant/config.py
@@ -735,18 +735,17 @@ def validate_config(cfg: dict) -> None:
 
     if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
         history = _get_path(cfg, "stage2_model.autoregressive.history")
-        if history != "markov":
+        if history not in ("markov", "attention"):
             raise ValueError(
-                f"stage2_model.autoregressive.history = {history!r} is "
-                "accepted by the schema but not implemented until v0.3.0 "
-                "step 7 — use 'markov'"
+                f"stage2_model.autoregressive.history = {history!r} — must "
+                "be 'markov' or 'attention'"
             )
         teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
-        if teacher_forcing != "always":
+        if teacher_forcing not in ("always", "scheduled", "never"):
             raise ValueError(
                 "stage2_model.autoregressive.teacher_forcing = "
-                f"{teacher_forcing!r} is accepted by the schema but not "
-                "implemented until v0.3.0 step 7 — use 'always'"
+                f"{teacher_forcing!r} — must be 'always', 'scheduled' or "
+                "'never'"
             )
 
 
diff --git a/giant/model/network.py b/giant/model/network.py
index 0624e13..a058af6 100644
--- a/giant/model/network.py
+++ b/giant/model/network.py
@@ -789,9 +789,16 @@ def build_trunk(
 
 class HistoryEncoder(nn.Module):
     """Interface for stage-2 autoregressive per-token history summaries:
-    `forward(feat, has_prev) -> (B, K, out_dim)`. `MarkovHistory` is the only
-    implementation until v0.3.0 step 7 (`AttentionHistory`,
-    `history = "attention"`)."""
+    `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
+    a full (teacher-forced) token sequence — used by training. `MarkovHistory`
+    and `AttentionHistory` (docs/v0.3.0-design.md §6.2) are the two
+    implementations. Inference (`giant/sample.py`) generates one token at a
+    time and cannot afford `forward`'s per-step cost to be O(K) (attention
+    would then be O(K^2) over a rollout's k_max loop); encoders that need
+    incremental state for that path additionally implement `init_cache`/
+    `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
+    its per-step cost is already O(1) (it only ever looks at the previous
+    token, not the full prefix)."""
 
     def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
@@ -819,6 +826,126 @@ class MarkovHistory(HistoryEncoder):
         return self.mlp(x)
 
 
+class _CausalAttnBlock(nn.Module):
+    """One pre-norm causal self-attention block for `AttentionHistory`.
+
+    Exposes two forward paths that must agree (see
+    `test_attention_history_step_matches_forward` in `tests/test_network.py`):
+    `forward` — the full-sequence, causally-masked pass used for training;
+    `step` — an incremental pass for inference, given the *pre-attention*
+    normalized hidden states of every earlier position (`kv_cache`, i.e.
+    `norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather
+    than raw `x` is what makes `step` correct: this block's attention needs
+    exactly that quantity as keys/values, and `LayerNorm` has no cross-position
+    interaction, so recomputing it per position instead of caching it would
+    still be correct but pointlessly repeat work. The *next* block's cache is
+    built from a different sequence (this block's output), so each block owns
+    an independent cache entry.
+    """
+
+    def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(
+            dim, n_heads, dropout=dropout, batch_first=True
+        )
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
+        )
+
+    def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
+        x = x + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    def step(
+        self, x_new: torch.Tensor, kv_cache: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
+        (first position) or `(B, T, dim)` — `norm1(x)` of every earlier
+        position at this same block. Returns `(out, new_kv_cache)`, `out`
+        being this position's block output (`(B, 1, dim)`, to feed the next
+        block's `step`), `new_kv_cache` the same cache extended by this
+        position (to reuse at this block's *next* `step` call)."""
+        h_new = self.norm1(x_new)
+        kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1)
+        attn_out, _ = self.attn(h_new, kv, kv, need_weights=False)
+        x = x_new + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x, kv
+
+
+class AttentionHistory(HistoryEncoder):
+    """Causal self-attention over the emitted-token prefix
+    (docs/v0.3.0-design.md §6.2) — the more expressive alternative to
+    `MarkovHistory`'s fixed previous-token-only summary. `feat`/`has_prev`
+    follow the same shifted-by-one convention `MarkovHistory` and
+    `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
+    own `(energy_fraction, direction, type_representation)`, with a learned
+    start vector substituted at `has_prev == False` positions (only slot 0 in
+    practice — see `giant.train._ar_has_prev`). Causal masking then makes
+    position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens
+    `0..i-1` — exactly the prefix available when predicting token `i`.
+
+    `forward` is the parallel training path (one pass over the whole
+    teacher-forced sequence); `init_cache`/`step` are the incremental
+    inference path `giant/sample.py` uses, one new token per call, to avoid
+    re-encoding the whole prefix from scratch every slot (docs/v0.3.0-design.md
+    §10's "KV cache" note) — `step` must be called exactly once per slot (its
+    cache-extension is not idempotent), so a slot's output must be reused for
+    every model call within that slot (`forward`'s ODE substeps, or a separate
+    `predict_type` call) rather than re-derived — see
+    `Stage2Autoregressive.history_step`.
+    """
+
+    def __init__(
+        self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2
+    ) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.in_proj = nn.Linear(in_dim, out_dim)
+        self.blocks = nn.ModuleList(
+            [_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]
+        )
+
+    def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.in_proj(x)
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        x = self._embed(feat, has_prev)
+        mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
+        for block in self.blocks:
+            x = block(x, mask)
+        return x
+
+    def init_cache(self) -> list[torch.Tensor | None]:
+        return [None for _ in self.blocks]
+
+    def step(
+        self,
+        token_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        cache: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
+        token's own features (what would be `feat[:, k]` in `forward`).
+        Advances every block's cache by this position and returns this
+        position's output (`(B, 1, out_dim)`, the correct history summary for
+        the NEXT slot) plus the updated cache."""
+        x = self._embed(token_feat, has_prev)
+        new_cache: list[torch.Tensor | None] = []
+        for block, kv in zip(self.blocks, cache):
+            x, kv_new = block.step(x, kv)
+            new_cache.append(kv_new)
+        return x, new_cache
+
+
 # ---------------------------------------------------------------------------
 # Stage models (docs/v0.3.0-design.md §5.3)
 # ---------------------------------------------------------------------------
@@ -1082,8 +1209,8 @@ class Stage2OneShot(nn.Module):
 class Stage2Autoregressive(nn.Module):
     """Emits secondaries one at a time in descending-energy order
     (docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous
-    k_max-slot prediction. Only `history = "markov"` is implemented (v0.3.0
-    step 5) — `history = "attention"` raises immediately at construction.
+    k_max-slot prediction. `history` selects `MarkovHistory` or
+    `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
     `teacher_forcing` handling lives entirely in the trainer
     (`giant/train.py`), since it only affects how training inputs are
     assembled, not this module's architecture.
@@ -1122,13 +1249,16 @@ class Stage2Autoregressive(nn.Module):
         build_n_sec_head: bool = True,
         particle_type_cfg: dict | None = None,
         history: str = "markov",
+        attn_n_heads: int = 4,
+        attn_n_layers: int = 2,
     ) -> None:
         super().__init__()
-        if history != "markov":
-            raise NotImplementedError(
-                f"stage2_model.autoregressive.history={history!r} is not "
-                "implemented until v0.3.0 step 7 — use 'markov'"
+        if history not in ("markov", "attention"):
+            raise ValueError(
+                f"stage2_model.autoregressive.history={history!r} — must be "
+                "'markov' or 'attention'"
             )
+        self.history_kind = history
         self.generator_kind = generator
         self.noise_dim = noise_dim
         self.k_max = k_max
@@ -1149,7 +1279,14 @@ class Stage2Autoregressive(nn.Module):
         # width — there's no dedicated stage2_model.autoregressive key for
         # this, a reasonable default rather than a design-doc-specified value.
         history_dim = cond_out_dim
-        self.history_encoder = MarkovHistory(CONT_SLOT_DIM + self.type_dim, history_dim)
+        hist_in_dim = CONT_SLOT_DIM + self.type_dim
+        self.history_encoder: HistoryEncoder = (
+            AttentionHistory(
+                hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
+            )
+            if history == "attention"
+            else MarkovHistory(hist_in_dim, history_dim)
+        )
         token_fuse_in = (
             cond_out_dim + context_dim + history_dim + 2
         )  # +2: remaining_frac, slot_idx
@@ -1205,14 +1342,48 @@ class Stage2Autoregressive(nn.Module):
         has_prev: torch.Tensor,
         remaining_frac: torch.Tensor,
         slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        """`hist`, if given, overrides recomputing `self.history_encoder`
+        from `history_feat`/`has_prev` — the inference-time KV-cache path
+        (`Stage2Autoregressive.history_step`) precomputes it once per slot and
+        passes it in here so a slot's (possibly several) model calls — an ODE
+        loop's substeps, or a separate `predict_type` call — read the same
+        cached history instead of each re-deriving (and, under attention,
+        re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
         K = history_feat.size(1)
         base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
         ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
-        hist = self.history_encoder(history_feat, has_prev)
+        if hist is None:
+            hist = self.history_encoder(history_feat, has_prev)
         scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
         return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
 
+    def init_history_cache(self):
+        """Inference-only incremental-decoding state for `self.history_encoder`
+        (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its
+        per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
+        `AttentionHistory.init_cache()` under `history="attention"`."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.init_cache()
+        return None
+
+    def history_step(
+        self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache
+    ) -> tuple[torch.Tensor, object]:
+        """One inference slot's worth of history encoding: advances `cache`
+        (from `init_history_cache`, or a previous `history_step` call) by
+        `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
+        token, same convention `giant.sample.sample_secondaries_ar` already
+        threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
+        this slot's history summary (pass it as `_token_cond`'s `hist=` to
+        every model call made for this slot), `new_cache` is what to pass into
+        the *next* slot's `history_step`. Must be called exactly once per
+        slot — see `AttentionHistory.step`'s docstring."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.step(token_feat, has_prev, cache)
+        return self.history_encoder(token_feat, has_prev), cache
+
     def forward(
         self,
         x_t: torch.Tensor,
@@ -1224,6 +1395,7 @@ class Stage2Autoregressive(nn.Module):
         remaining_frac: torch.Tensor,
         slot_idx: torch.Tensor,
         t: torch.Tensor | None = None,
+        hist: torch.Tensor | None = None,
     ) -> torch.Tensor:
         B, K = x_t.shape[0], x_t.shape[1]
         c_emb = self._token_cond(
@@ -1234,6 +1406,7 @@ class Stage2Autoregressive(nn.Module):
             has_prev,
             remaining_frac,
             slot_idx,
+            hist=hist,
         )
         if self.time_emb is not None:
             assert t is not None
@@ -1268,6 +1441,7 @@ class Stage2Autoregressive(nn.Module):
         has_prev: torch.Tensor,
         remaining_frac: torch.Tensor,
         slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if self.type_head is None:
             raise RuntimeError(
@@ -1285,6 +1459,7 @@ class Stage2Autoregressive(nn.Module):
             has_prev,
             remaining_frac,
             slot_idx,
+            hist=hist,
         )
         B, K, _ = c_emb.shape
         return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
@@ -1571,6 +1746,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 build_n_sec_head=legacy_owner != "stage1",
                 particle_type_cfg=particle_type_cfg,
                 history=ar_cfg.get("history", "markov"),
+                attn_n_heads=ar_cfg.get("attn_n_heads", 4),
+                attn_n_layers=ar_cfg.get("attn_n_layers", 2),
             )
         else:
             sec_dim = stage2_trunk_sec_dim(
diff --git a/giant/sample.py b/giant/sample.py
index f9ba019..2f8df2f 100644
--- a/giant/sample.py
+++ b/giant/sample.py
@@ -256,17 +256,28 @@ def sample_secondaries_ar(
     once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
     physics step.
 
+    Under `history="attention"` the history encoding is computed once per
+    slot via `Stage2Autoregressive.history_step` (a KV-cache append, §10)
+    rather than re-derived by every model call inside that slot — so an ODE
+    loop's `steps` substeps, and the separate `predict_type` call when the
+    type slice isn't folded into the trunk output, all reuse the SAME `hist`
+    tensor for a given `k`. Recomputing per call instead would be merely
+    wasteful under markov (its per-call cost is already O(1)) but wrong under
+    attention: `AttentionHistory.step` mutates the cache by appending, so
+    calling it more than once per slot would double-count that slot's own
+    (not-yet-existing) predecessor.
+
     The free-running history feature stays UNSNAPPED (mirrors the
     established "no snapping" precedent for `particle_type.target =
     "physical"` secondaries feeding their own future conditioning):
     `"physical"` carries the raw (log_mass, charge) forward as-is;
     `"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
     one exception — its history slot must be a probability-simplex-shaped
-    vector (that's what `MarkovHistory` was trained on, `_type_repr`'s
-    `F.one_hot` ground truth), so it's the hard one-hot of `argmax(logits)`,
-    not the raw logits themselves. Discretizing further, into a concrete PDG
-    code, only ever happens once — at secondary-spawn time in
-    `giant/rollout.py` — never inside this loop.
+    vector (that's what `MarkovHistory`/`AttentionHistory` were trained on,
+    `_type_repr`'s `F.one_hot` ground truth), so it's the hard one-hot of
+    `argmax(logits)`, not the raw logits themselves. Discretizing further,
+    into a concrete PDG code, only ever happens once — at secondary-spawn
+    time in `giant/rollout.py` — never inside this loop.
 
     Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
     `sample_secondaries`/`sample_secondaries_wgan`'s (see
@@ -291,6 +302,7 @@ def sample_secondaries_ar(
     # Running per-token state, threaded from one slot to the next.
     prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
     remaining = torch.ones(B, device=device)
+    history_cache = sec_decoder.init_history_cache()
 
     for k in range(k_max):
         has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
@@ -299,6 +311,9 @@ def sample_secondaries_ar(
         slot_idx = torch.full(
             (B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
         )
+        hist, history_cache = sec_decoder.history_step(
+            history_feat, has_prev, history_cache
+        )
 
         if generator == "wgan":
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -311,6 +326,7 @@ def sample_secondaries_ar(
                 has_prev,
                 remaining_frac,
                 slot_idx,
+                hist=hist,
             )
         else:
             x = torch.randn(B, 1, token_dim, device=device)
@@ -327,6 +343,7 @@ def sample_secondaries_ar(
                     remaining_frac,
                     slot_idx,
                     t=t,
+                    hist=hist,
                 )
                 x = x + v * dt
             token = x
@@ -344,6 +361,7 @@ def sample_secondaries_ar(
                 has_prev,
                 remaining_frac,
                 slot_idx,
+                hist=hist,
             ).squeeze(1)
 
         sec_cont[:, k] = cont_k
diff --git a/giant/train.py b/giant/train.py
index a504056..6a3689c 100644
--- a/giant/train.py
+++ b/giant/train.py
@@ -27,6 +27,7 @@ from giant.model.schedule import (
     flow_matching_loss_secondary_ar,
 )
 from giant.model.wgan import gradient_penalty, generator_loss
+from giant.sample import sample_secondaries_ar
 from giant.validate import validate_marginals
 
 _CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
@@ -246,6 +247,119 @@ def _assemble_stage2_ar_inputs(
     }
 
 
+def _stage2_tf_prob(
+    mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
+) -> float:
+    """P(condition slot k+1 on the TRUE token k rather than the model's own
+    prediction), for the current epoch (docs/v0.3.0-design.md §3.3
+    `stage2_model.autoregressive.teacher_forcing`). `"always"`/`"never"` are
+    the two degenerate constants; `"scheduled"` linearly interpolates
+    `p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled
+    sampling (Bengio et al. 2015)."""
+    if mode == "always":
+        return 1.0
+    if mode == "never":
+        return 0.0
+    frac = epoch / max(total_epochs - 1, 1)
+    frac = min(max(frac, 0.0), 1.0)
+    return p_start + (p_end - p_start) * frac
+
+
+def _history_repr_from_ar_sample(
+    sec_cont_pred: torch.Tensor,
+    sec_type_pred: torch.Tensor,
+    particle_type_cfg: dict,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """`(fraction, direction, type_repr)` — the same triple `_type_repr` /
+    `_stick_fraction` derive from ground truth, but from a free-running
+    `sample_secondaries_ar` self-sample instead, so the two can be mixed
+    slot-by-slot under scheduled sampling (`_assemble_stage2_ar_inputs_scheduled`).
+    `target="onehot"` collapses the raw per-slot type logits to a hard
+    one-hot of `argmax` — `sample_secondaries_ar`'s own history convention
+    (see its docstring), matching what `MarkovHistory`/`AttentionHistory`
+    were trained on; the other two targets are already the right
+    representation."""
+    fraction = torch.sigmoid(sec_cont_pred[..., 0])
+    direction = sec_cont_pred[..., 1:4]
+    if particle_type_cfg.get("target", "physical") == "onehot":
+        type_dim = sec_type_pred.size(-1)
+        type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float()
+    else:
+        type_repr = sec_type_pred
+    return fraction, direction, type_repr
+
+
+def _assemble_stage2_ar_inputs_scheduled(
+    model: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_ctx: torch.Tensor,
+    sec_cont: torch.Tensor,
+    sec_type_idx: torch.Tensor,
+    n_sec: torch.Tensor,
+    particle_type_cfg: dict,
+    cond_enc: torch.nn.Module,
+    emb_dim: int,
+    p_tf: float,
+    sample_steps: int,
+) -> dict[str, torch.Tensor]:
+    """Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs`
+    (docs/v0.3.0-design.md §3.3 `teacher_forcing` = "scheduled"/"never"):
+    each slot's history is the TRUE previous token with probability `p_tf`
+    (an independent per-example, per-slot Bernoulli draw) and the model's own
+    free-running prediction otherwise — closing the train/inference gap that
+    `teacher_forcing="always"` (ground truth throughout training) never sees.
+    `p_tf >= 1.0` degenerates exactly to `_assemble_stage2_ar_inputs` (and
+    skips self-sampling entirely), so callers can call this unconditionally.
+
+    The free-running estimate is a REAL autoregressive self-sample —
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
+    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
+    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    this is called on (§6.4's cost note, paid at train time too whenever
+    teacher_forcing != "always"). Fully detached: gradient only ever flows
+    through the "real" target path each stage trainer already uses
+    (`_assemble_stage2_ar_target`), never through this self-sample.
+    """
+    device = sec_cont.device
+    B, K = sec_cont.shape[0], sec_cont.shape[1]
+    if p_tf >= 1.0:
+        return _assemble_stage2_ar_inputs(
+            sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim
+        )
+
+    was_training = model.training
+    sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
+        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
+    )
+    if was_training:
+        model.train()
+
+    fraction_gt = _stick_fraction(sec_cont)
+    dir_gt = sec_cont[..., 1:4]
+    type_repr_gt = _type_repr(
+        sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim
+    )
+    fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
+        sec_cont_pred, sec_type_pred, particle_type_cfg
+    )
+
+    use_gt = torch.rand(B, K, device=device) < p_tf
+    fraction = torch.where(use_gt, fraction_gt, fraction_pred)
+    direction = torch.where(use_gt.unsqueeze(-1), dir_gt, dir_pred)
+    type_repr = torch.where(use_gt.unsqueeze(-1), type_repr_gt, type_repr_pred)
+
+    own_feat = torch.cat([fraction.unsqueeze(-1), direction, type_repr], dim=-1)
+    return {
+        "history_feat": _shift_prev(own_feat),
+        "has_prev": _ar_has_prev(K, device).expand(B, -1),
+        "remaining_frac": _remaining_energy_fraction(fraction),
+        "slot_idx": (torch.arange(K, device=device).float() / max(K - 1, 1))
+        .unsqueeze(0)
+        .expand(B, -1),
+    }
+
+
 def _relax_onehot_type_slice(
     x_flat: torch.Tensor,
     k_max: int,
@@ -362,6 +476,10 @@ class FlowDDPMStageTrainer(StageTrainer):
         particle_type_cfg: dict | None = None,
         particle_type_emb_dim: int = 16,
         decoder: str = "one_shot",
+        teacher_forcing: str = "always",
+        tf_p_start: float = 1.0,
+        tf_p_end: float = 1.0,
+        ar_sample_steps: int = 10,
     ) -> None:
         if is_stage2 and generator not in ("flow",):
             raise NotImplementedError(
@@ -373,6 +491,12 @@ class FlowDDPMStageTrainer(StageTrainer):
         self.is_stage2 = is_stage2
         self.generator = generator
         self.decoder = decoder
+        self.teacher_forcing = teacher_forcing
+        self.tf_p_start = tf_p_start
+        self.tf_p_end = tf_p_end
+        self.ar_sample_steps = ar_sample_steps
+        self.total_epochs = epochs
+        self.steps_per_epoch = max(steps_per_epoch, 1)
         self.device = device
         self.model = model.to(device)
         self.lambda_weight = lambda_weight
@@ -513,7 +637,13 @@ class FlowDDPMStageTrainer(StageTrainer):
             l_type = (se * mask).sum() / denom
         return l_type, type_acc
 
-    def _compute(self, batch: tuple, device: torch.device) -> dict:
+    def _compute(
+        self, batch: tuple, device: torch.device, epoch: int | None = None
+    ) -> dict:
+        """`epoch=None` (the `val_loss` path) always uses full teacher
+        forcing (`p_tf=1.0`) regardless of `self.teacher_forcing` — validation
+        should stay a stable, non-stochastic ground-truth comparison; only
+        the training `step` path schedules `p_tf` by epoch."""
         (
             cond_cont,
             cond_cat,
@@ -529,12 +659,30 @@ class FlowDDPMStageTrainer(StageTrainer):
         x1_s2 = None
         ar_inputs = None
         if self.is_stage2 and self.decoder == "autoregressive":
-            ar_inputs = _assemble_stage2_ar_inputs(
+            p_tf = (
+                1.0
+                if epoch is None
+                else _stage2_tf_prob(
+                    self.teacher_forcing,
+                    self.tf_p_start,
+                    self.tf_p_end,
+                    epoch,
+                    self.total_epochs,
+                )
+            )
+            ar_inputs = _assemble_stage2_ar_inputs_scheduled(
+                self.model,
+                cond_cont,
+                cond_cat,
+                stage1_ctx,
                 sec_cont,
                 sec_type_idx,
+                n_sec,
                 self.particle_type_cfg,
                 self.model.cond_enc,
                 self.particle_type_emb_dim,
+                p_tf,
+                self.ar_sample_steps,
             )
             x1_s2 = _assemble_stage2_ar_target(
                 sec_cont,
@@ -613,7 +761,8 @@ class FlowDDPMStageTrainer(StageTrainer):
                 self.gumbel_tau_start,
                 self.gumbel_tau_end,
             )
-        out = self._compute(batch, device)
+        epoch = global_step // self.steps_per_epoch
+        out = self._compute(batch, device, epoch=epoch)
         self.optimizer.zero_grad()
         out["total"].backward()
         grad_norm = torch.nn.utils.clip_grad_norm_(self.params, 1.0)
@@ -719,10 +868,20 @@ class WGANStageTrainer(StageTrainer):
         type_gumbel_tau_start: float = 1.0,
         type_gumbel_tau_end: float = 0.1,
         decoder: str = "one_shot",
+        teacher_forcing: str = "always",
+        tf_p_start: float = 1.0,
+        tf_p_end: float = 1.0,
+        ar_sample_steps: int = 10,
     ) -> None:
         self.name = name
         self.is_stage2 = is_stage2
         self.decoder = decoder
+        self.teacher_forcing = teacher_forcing
+        self.tf_p_start = tf_p_start
+        self.tf_p_end = tf_p_end
+        self.ar_sample_steps = ar_sample_steps
+        self.total_epochs = epochs
+        self.steps_per_epoch = max(steps_per_epoch, 1)
         self.device = device
         self.model = model.to(device)
         self.critic = critic.to(device)
@@ -810,12 +969,27 @@ class WGANStageTrainer(StageTrainer):
                 return self.critic(x, cond_cont, cond_cat, stage1_ctx)
 
             if self.decoder == "autoregressive":
-                ar = _assemble_stage2_ar_inputs(
+                epoch = global_step // self.steps_per_epoch
+                p_tf = _stage2_tf_prob(
+                    self.teacher_forcing,
+                    self.tf_p_start,
+                    self.tf_p_end,
+                    epoch,
+                    self.total_epochs,
+                )
+                ar = _assemble_stage2_ar_inputs_scheduled(
+                    self.model,
+                    cond_cont,
+                    cond_cat,
+                    stage1_ctx,
                     sec_cont,
                     sec_type_idx,
+                    n_sec,
                     self.particle_type_cfg,
                     self.model.cond_enc,
                     self.particle_type_emb_dim,
+                    p_tf,
+                    self.ar_sample_steps,
                 )
                 real = (
                     _assemble_stage2_ar_target(
@@ -1051,16 +1225,15 @@ def _build_stage_trainers(
         }
         particle_type_emb_dim = cfg["conditioning"]["particle"]["emb_dim"]
         decoder = stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot"
-        if is_stage2 and decoder == "autoregressive":
-            teacher_forcing = (stage_cfg.get("autoregressive") or {}).get(
-                "teacher_forcing", "always"
-            )
-            if teacher_forcing != "always":
-                raise NotImplementedError(
-                    "stage2_model.autoregressive.teacher_forcing="
-                    f"{teacher_forcing!r} is not implemented until v0.3.0 "
-                    "step 7 — use 'always'"
-                )
+        ar_cfg = (stage_cfg.get("autoregressive") or {}) if is_stage2 else {}
+        teacher_forcing = ar_cfg.get("teacher_forcing", "always")
+        tf_p_start = ar_cfg.get("tf_p_start", 1.0)
+        tf_p_end = ar_cfg.get("tf_p_end", 1.0)
+        # AR self-sampling under scheduled/never teacher forcing reuses
+        # train.validate_steps as its flow-matching ODE step count — no
+        # dedicated config key for this (docs/v0.3.0-design.md §3.3 lists
+        # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
+        ar_sample_steps = t.get("validate_steps", 10)
 
         if generator == "wgan":
             critic = critics.get(name)
@@ -1090,6 +1263,10 @@ def _build_stage_trainers(
                 type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
                 type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
                 decoder=decoder,
+                teacher_forcing=teacher_forcing,
+                tf_p_start=tf_p_start,
+                tf_p_end=tf_p_end,
+                ar_sample_steps=ar_sample_steps,
             )
         else:
             ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000)
@@ -1116,6 +1293,10 @@ def _build_stage_trainers(
                 particle_type_cfg=particle_type_cfg,
                 particle_type_emb_dim=particle_type_emb_dim,
                 decoder=decoder,
+                teacher_forcing=teacher_forcing,
+                tf_p_start=tf_p_start,
+                tf_p_end=tf_p_end,
+                ar_sample_steps=ar_sample_steps,
             )
     return trainers
 
diff --git a/tests/test_config.py b/tests/test_config.py
index 4fb9b00..b57b5cf 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -1,6 +1,8 @@
 from datetime import datetime
 from pathlib import Path
 
+import pytest
+
 from giant import config as gconfig
 
 _CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"
@@ -645,13 +647,37 @@ def test_validate_config_ar_default_markov_always_passes():
     gconfig.validate_config(cfg)  # must not raise
 
 
-def test_validate_config_ar_history_attention_not_implemented():
+def test_validate_config_ar_history_attention_passes():
+    """v0.3.0 step 7 implements history='attention' — must not raise."""
     cfg = _cfg_with(
         **{
             "stage2_model.decoder": "autoregressive",
             "stage2_model.autoregressive.history": "attention",
         }
     )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+@pytest.mark.parametrize("teacher_forcing", ["scheduled", "never"])
+def test_validate_config_ar_teacher_forcing_scheduled_or_never_passes(teacher_forcing):
+    """v0.3.0 step 7 implements teacher_forcing in {'scheduled', 'never'} —
+    must not raise."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.teacher_forcing": teacher_forcing,
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_ar_history_invalid_value_rejected():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.history": "bogus",
+        }
+    )
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
@@ -659,11 +685,11 @@ def test_validate_config_ar_history_attention_not_implemented():
         assert "history" in str(e)
 
 
-def test_validate_config_ar_teacher_forcing_scheduled_not_implemented():
+def test_validate_config_ar_teacher_forcing_invalid_value_rejected():
     cfg = _cfg_with(
         **{
             "stage2_model.decoder": "autoregressive",
-            "stage2_model.autoregressive.teacher_forcing": "scheduled",
+            "stage2_model.autoregressive.teacher_forcing": "bogus",
         }
     )
     try:
diff --git a/tests/test_network.py b/tests/test_network.py
index dde4bf9..f3c5293 100644
--- a/tests/test_network.py
+++ b/tests/test_network.py
@@ -2,6 +2,7 @@ import pytest
 import torch
 from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
+    AttentionHistory,
     ConditionEncoder,
     MarkovHistory,
     SinusoidalEmbedding,
@@ -323,6 +324,70 @@ def test_markov_history_uses_start_vector_when_no_prev():
     assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
 
 
+# --- AttentionHistory (docs/v0.3.0-design.md §6.2, v0.3.0 step 7) ----------
+
+
+def test_attention_history_shape():
+    hist = AttentionHistory(in_dim=7, out_dim=12, n_heads=2, n_layers=2)
+    B, K = 3, 5
+    feat = torch.randn(B, K, 7)
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    out = hist(feat, has_prev)
+    assert out.shape == (B, K, 12)
+
+
+def test_attention_history_uses_start_vector_when_no_prev():
+    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=1)
+    B, K = 2, 3
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    feat_a = torch.randn(B, K, 4)
+    feat_b = feat_a.clone()
+    feat_b[:, 0] = torch.randn(B, 4) * 100
+    out_a = hist(feat_a, has_prev)
+    out_b = hist(feat_b, has_prev)
+    assert torch.allclose(out_a[:, 0], out_b[:, 0], atol=1e-5)
+
+
+def test_attention_history_is_causal():
+    """Position i's output must not depend on feat at positions > i — unlike
+    MarkovHistory (which only ever looks at position i itself, already
+    trivially "causal"), this is AttentionHistory's actual contribution:
+    seeing the full prefix 0..i-1, never anything later."""
+    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
+    hist.eval()
+    B, K = 2, 5
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    feat_a = torch.randn(B, K, 4)
+    feat_b = feat_a.clone()
+    feat_b[:, 3:] = torch.randn(B, K - 3, 4) * 100
+    with torch.no_grad():
+        out_a = hist(feat_a, has_prev)
+        out_b = hist(feat_b, has_prev)
+    assert torch.allclose(out_a[:, :3], out_b[:, :3], atol=1e-5)
+
+
+def test_attention_history_step_matches_forward():
+    """The incremental KV-cache path (`init_cache`/`step`,
+    `giant/sample.py`'s AR loop) must reproduce `forward`'s parallel-pass
+    output exactly, one position at a time."""
+    hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
+    hist.eval()
+    B, K = 3, 6
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    feat = torch.randn(B, K, 4)
+    with torch.no_grad():
+        expected = hist(feat, has_prev)
+
+        cache = hist.init_cache()
+        outs = []
+        for k in range(K):
+            out_k, cache = hist.step(feat[:, k : k + 1], has_prev[:, k : k + 1], cache)
+            outs.append(out_k)
+        stepped = torch.cat(outs, dim=1)
+
+    assert torch.allclose(stepped, expected, atol=1e-5)
+
+
 # --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) ---------
 
 
@@ -361,16 +426,19 @@ def _ar_inputs(B: int, K: int, hist_dim: int):
     return history_feat, has_prev, remaining_frac, slot_idx
 
 
-def test_stage2_autoregressive_history_attention_raises():
-    with pytest.raises(NotImplementedError):
-        _build_stage2_ar("onehot", "wgan", history="attention")
+def test_stage2_autoregressive_history_invalid_raises():
+    with pytest.raises(ValueError):
+        _build_stage2_ar("onehot", "wgan", history="bogus")
 
 
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
-def test_stage2_autoregressive_forward_shape(target, generator):
+@pytest.mark.parametrize("history", ["markov", "attention"])
+def test_stage2_autoregressive_forward_shape(target, generator, history):
     B, K, emb_dim = 4, 5, 6
-    model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
+    model = _build_stage2_ar(
+        target, generator, emb_dim=emb_dim, k_max=K, history=history
+    )
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
@@ -518,3 +586,45 @@ def test_stage2_autoregressive_gradients_flow_onehot():
     (flow_out + nsec_out + type_out).backward()
     for name, p in model.named_parameters():
         assert p.grad is not None, f"no grad for {name}"
+
+
+def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
+    """`init_history_cache`/`history_step` (the incremental path
+    `giant/sample.py`'s AR loop drives, one slot per call) must reproduce
+    exactly what one parallel `self.history_encoder(history_feat, has_prev)`
+    call over the whole shifted sequence would give at each position — the
+    KV-cache correctness guarantee, exercised through `Stage2Autoregressive`
+    itself rather than `AttentionHistory` in isolation
+    (`test_attention_history_step_matches_forward` covers that lower layer)."""
+    B, K, emb_dim = 3, 6, 6
+    model = _build_stage2_ar(
+        "physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
+    )
+    model.eval()
+    type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
+    hist_in_dim = CONT_SLOT_DIM + type_dim
+    own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
+    has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    history_feat = torch.cat(
+        [torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
+    )
+
+    with torch.no_grad():
+        expected = model.history_encoder(history_feat, has_prev_full)
+
+        cache = model.init_history_cache()
+        outs = []
+        prev = torch.zeros(B, 1, hist_in_dim)
+        for k in range(K):
+            has_prev_k = torch.full((B, 1), k >= 1, dtype=torch.bool)
+            hist_k, cache = model.history_step(prev, has_prev_k, cache)
+            outs.append(hist_k)
+            prev = own_feat[:, k : k + 1]
+        stepped = torch.cat(outs, dim=1)
+
+    assert torch.allclose(stepped, expected, atol=1e-5)
+
+
+def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
+    model = _build_stage2_ar("physical", "wgan", history="markov")
+    assert model.init_history_cache() is None
diff --git a/tests/test_sample.py b/tests/test_sample.py
index 670942a..cb8bdfd 100644
--- a/tests/test_sample.py
+++ b/tests/test_sample.py
@@ -76,6 +76,7 @@ def _stage2_ar(
     pdg: int = 3,
     mat: int = 2,
     k_max: int = 5,
+    history: str = "markov",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(
         _conditioning_for(target), emb_dim
@@ -92,6 +93,9 @@ def _stage2_ar(
         noise_dim=8,
         k_max=k_max,
         particle_type_cfg={"target": target},
+        history=history,
+        attn_n_heads=2,
+        attn_n_layers=1,
     ).eval()
 
 
@@ -187,11 +191,14 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
 # ── Stage2Autoregressive ─────────────────────────────────────────────────────
 
 
+@pytest.mark.parametrize("history", ["markov", "attention"])
 @pytest.mark.parametrize("generator", ["flow", "wgan"])
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
-def test_sample_secondaries_ar_shapes(target, generator):
+def test_sample_secondaries_ar_shapes(target, generator, history):
     B, k_max, emb_dim = 4, 5, 6
-    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
+    decoder = _stage2_ar(
+        target, generator, emb_dim=emb_dim, k_max=k_max, history=history
+    )
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, k_max + 1, (B,))
diff --git a/tests/test_train.py b/tests/test_train.py
index 8318add..3e37769 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -2,6 +2,7 @@
 
 import copy
 import csv
+import math
 import tempfile
 from pathlib import Path
 
@@ -29,6 +30,7 @@ from giant.train import (
     _relax_onehot_type_slice,
     _remaining_energy_fraction,
     _shift_prev,
+    _stage2_tf_prob,
     _stick_fraction,
     _type_repr,
     _wandb_run_config,
@@ -114,6 +116,37 @@ def test_ar_has_prev_false_only_at_slot_zero():
     assert has_prev.tolist() == [[False, True, True, True, True]]
 
 
+# --- _stage2_tf_prob (docs/v0.3.0-design.md §3.3, v0.3.0 step 7) -----------
+
+
+def test_stage2_tf_prob_always_is_constant_one():
+    assert _stage2_tf_prob("always", 1.0, 0.0, 0, 10) == 1.0
+    assert _stage2_tf_prob("always", 1.0, 0.0, 9, 10) == 1.0
+
+
+def test_stage2_tf_prob_never_is_constant_zero():
+    assert _stage2_tf_prob("never", 1.0, 1.0, 0, 10) == 0.0
+    assert _stage2_tf_prob("never", 1.0, 1.0, 9, 10) == 0.0
+
+
+def test_stage2_tf_prob_scheduled_interpolates_linearly():
+    assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 11) == 1.0
+    assert abs(_stage2_tf_prob("scheduled", 1.0, 0.0, 5, 11) - 0.5) < 1e-9
+    assert _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) == 0.0
+
+
+def test_stage2_tf_prob_scheduled_clamps_beyond_total_epochs():
+    end = _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11)
+    beyond = _stage2_tf_prob("scheduled", 1.0, 0.0, 50, 11)
+    assert beyond == end
+
+
+def test_stage2_tf_prob_scheduled_handles_single_epoch():
+    # total_epochs=1 is guarded to a denominator of 1 internally (like
+    # _gumbel_tau's total_steps=0 guard) — epoch=0 gives zero progress.
+    assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 1) == 1.0
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 def test_type_repr_shapes_and_values(target):
     B, K, emb_dim = 3, 4, 6
@@ -592,20 +625,70 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
 # --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
 
 
-def test_build_stage_trainers_rejects_scheduled_teacher_forcing():
+@pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
+@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
+    teacher_forcing, history, stage2_generator
+):
+    """v0.3.0 step 7: history='attention' and teacher_forcing in
+    {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
+    must run and produce a finite loss, for every {history} x
+    {teacher_forcing} x {generator} combination."""
     cfg = _base_cfg()
     cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["generator"] = stage2_generator
     cfg["stage2_model"]["autoregressive"] = {
-        "history": "markov",
-        "teacher_forcing": "scheduled",
+        "history": history,
+        "teacher_forcing": teacher_forcing,
+        "tf_p_start": 1.0,
+        "tf_p_end": 0.0,
+        "attn_n_heads": 2,
+        "attn_n_layers": 1,
     }
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
-    with pytest.raises(NotImplementedError):
-        _build_stage_trainers(
-            cfg, models, critics, torch.device("cpu"), total_train_batches=4
+    trainers = _build_stage_trainers(
+        cfg, models, critics, torch.device("cpu"), total_train_batches=4
+    )
+    trainer = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
+    loss_key = "g_loss" if stage2_generator == "wgan" else "loss"
+    assert math.isfinite(stats[loss_key])
+
+
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
+    stage2_generator,
+):
+    """Full `train()` run (not just one `trainer.step()` call) with
+    history='attention' AND teacher_forcing='scheduled' together — the
+    combination v0.3.0 step 7 exists to land — must complete and write a
+    checkpoint + metrics.csv with finite losses throughout."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["decoder"] = "autoregressive"
+    cfg["stage2_model"]["generator"] = stage2_generator
+    cfg["stage2_model"]["autoregressive"] = {
+        "history": "attention",
+        "teacher_forcing": "scheduled",
+        "tf_p_start": 1.0,
+        "tf_p_end": 0.0,
+        "attn_n_heads": 2,
+        "attn_n_layers": 1,
+    }
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert len(rows) == cfg["train"]["epochs"]
+        loss_col = (
+            "stage2_train_g_loss" if stage2_generator == "wgan" else "stage2_train_loss"
         )
+        assert all(math.isfinite(float(r[loss_col])) for r in rows)
 
 
 def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():